$\boldsymbol{\mu}\mathbf{P^2}$: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling
Authors: Moritz Haas
Jin Xu
Volkan Cevher
Leena Chennuru Vankadara
Venue: NeurIPS 2024 poster (2024)
Keywords: Deep Learning Theory
Optimal Hyperparameter Transfer
Sharpness Aware Minimization
Infinite Width Limits
Signal Propagation Theory
Tensor Programs
Primary area: optimization_for_deep_networks

Abstract: Sharpness Aware Minimization (SAM) enhances performance across various neural architectures and datasets. As models are continually scaled up to improve performance, a rigorous understanding of SAM’s scaling behaviour is paramount. To this end, we study the infinite-width limit of neural networks trained with SAM, using the Tensor Programs framework. Our findings reveal that the dynamics of standard SAM effectively reduce to applying SAM solely in the last layer in wide neural networks, even with optimal hyperparameters. In contrast, we identify a stable parameterization with layerwise perturbation scaling, which we call *Maximal Update and Perturbation Parameterization* ($\mu$P$^2$), that ensures all layers are both feature learning and effectively perturbed in the limit. Through experiments with MLPs, ResNets and Vision Transformers, we empirically demonstrate that $\mu$P$^2$ is the first parameterization to achieve hyperparameter transfer of the joint optimum of learning rate and perturbation radius across model scales. Moreover, we provide an intuitive condition to derive $\mu$P$^2$ for other perturbation rules like Adaptive SAM and SAM-ON, also ensuring balanced perturbation effects across all layers.

μP?: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling
Moritz Haas'
Jin Xu²
Volkan Cevher
Leena Chennuru Vankadara*
1 University of Tübingen, Tübingen Al Center" LIONS, EPFL"
University of Oxford" AGI Foundations, Amazon
Abstract
Sharpness Aware Minimization (SAM) enhances performance across various neural architectures and datasets, As models are continually scaled up to improve performance, rigorous understanding of SAM's scaling behaviour is paramount. To this end, we study the infinite-width limit of neural networks trained with SAM, using the Tensor Programs framework. Our findings reveal that the dynamics of standard SAM effectively reduce to applying SAM solely in the last layer in wide neural networks, even with optimal hyperparameters. In contrast, we identify a stable parameterization with layerwise perturbation scaling, which we call Maximal Update and Perturbution Parameterization (μP²), that ensures all layers are both feature learning and effectively perturbed in the limit. Through experiments with MLPs, ResNets and Vision Transformers, demonstrate that μP? achieves hyperparameter transfer of the joint optimum learning rate and perturbation radius across model scales. Moreover, we provide an intuitive condition to derive ju2 for other perturbation rules like Adaptive SAM and SAM-ON, also ensuring balanced perturbation effects across all layers.
empirically
1 Introduction
Sharpness Aware Minimization (SAM) (Foret et al., 2021) and its variants (Kwon al., 2021; Müller et al., 2024) improves generalization performance across wide range of neural architectures and datasets (Chen et al., 2021; Kaddour al., 2022). the SAM formulation, we minimize given loss L between our prediction and the data adversary simultaneously maximizes the same loss by perturbing the weights within a budget p.
function of the architecture's weights W, where an
A standard SAM update for an L-hidden layer multi layer perceptron (MLP) is given by
W W -m V WiL C (f (&s; W, + +x), 3)), with ef = p Vwi Wt). ve) W.). v:)l (SAM)
where W [W] WT+11, I1, t is the iteration count and E denotes the perturbation in the 1-th MLP layer with width n E N, and where we define an L-hidden layer MLP iteratively via
:= =- W'g.
x'(E) (E),
k++ = W (())
f(E) - - (E),
for inputs & E udd with trainable weight matrices W E gnd.... W guxa for L [2, L}. and WL+1 E Xn, We call h' preactivations, z' activations, and f({) output function. Despite the inherent difficulty of non-convex, non-concave optimization, SAM is quite successful in practice.
"This work was conducted during Moritz', Jin's and Volkan's time at Amazon. Correspondence to: mo.haas @ uni-tuebingen.de
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

IP P-global
frP2
width
Learning rale 1y
4096
1024
256
APP global
Perturbation radius P
Perturbation radius P
Perturbation radius
Left and center (μPP² transfers both and p): Test accuracy as function of learning rate 1j and perturbation radius p of a 3-layer MI.P in juP trained with SAM on CIFARIO for various widths with global perturbation scaling p 11/2 (left) and our layerwise perturbations scaling jPP (right), averaged over independent runs. 'x" denotes the optimum. Blue contours (the darker, the wider) denote the region within 1% of the optimal test accuracy smoothened with a Gaussian filter. Grey regions (the lighter. the wider) denote the unstable regime below 30% test accuracy. Right (NP² improves generalization): Same left but sliced at the optimal learning rate of both parameterizations for width 4096 with the base optimizer SGD in jP (dashed line) as baseline. Average and 2o-CI from 16 independent runs. Global perturbation scaling n-122 achieves width-indeepedent critical perturbation radius which training becomes unstable, but does not consistently improve over SGD in pP and does not transfer the optimal (1) p). jaP² achieves joint transfer in (1). p) and improves generalization performance.
Figure 1:
On the other hand, the steadily growing scale of foundation models has sparked considerable interest in scaling laws of model size and dataset size (Kaplan et al., 2020; Zhai et al., 2022). To rigorously understand learning dynamics under width scaling, Yang and Hu (2021) have recently provided general infinite-width theory for SGD, which has since been shown to be a good model for understanding the properties of large models (Vyas et al., 2024). Yang and Hu (2021) show that standard parameterizations (SP), including He or LeCun initialization (He et al.. 2015; LeCun et al., 2002) with a global learning rate, do not learn features the infinite-width limit.
Instead, a different scaling of layerwise initialization variances and learning rates, termed Maximal Update Purameterization (μPP), is necessary to achieve feature learning in all layers in wide networks. A crucial practical benefit of μP is the transfcrability of the optimal learning rate across model scales (Yang et al., 2022). This can drastically reduce computational costs as it allows to tune hyperparameters on smaller representative models and then to train the large model only once.
Contributions. In this paper, we adopt a scaling perspective to understand SAM's learning dynamics. Using the Tensor Programs framework (Yang, 2019; Yang and Hu, 2021; Yang and Littwin, 2023), this work provides the first infinite-width theory for SAM with important practical consequences:
1. We show that training an MLP with the standard (SAM) update rule is equivalent to applying perturbations only in the last layer the infinite- width limit, even if the perturbation radius is properly tuned. This holds for any width-dependent scaling of layerwise initialization variances and learning rates, including SP and jPP
2. We demonstrate that the optimal perturbation radius can shift significantly in jPP (Figure 1).
3. We postulate that jointly transferring the optimal learning rate 1j and perturbation radius P requires intth-independent feature learning and effective perturbations in every layer in the infinite- width limit. We define the perturbation of a trainable weight tensor to be effective iff its effect on the output function scales width hidth-independently. We show that this can be achieved with layerwise scalings of the perturbation radius, and provide a complete characterization of perturbation scaling parameterizations into four regimes: unstable, vanishing, nontrivial and effective perturbations.
4. We derive the Maximal Update and Perturbation Parameterization (nP2) that achieves both feature learning and effective perturbations in all layers in the infinite-width limit. We empirically demonstrate that juP2 achieves hyperparameter transfer in both learning rate n and perturbation radius (Figure 1).
5. We provide versatile (spectral) scaling condition (+) applicable to architectures such as ResNets and Vision Transformers (ViTs), and to various SAM variants like SAM-ON and Adaptive SAM (ASAM), and any SAM updates modeled Tensor Program.

---

2 Background and related work
We here provide short summary of related work. A more detailed account is provided in Appendix B. Sharpness Aware Minimization. SAM was motivated an inductive bias towards flatter minima and it provably reduces properties of the Hessian that are related to sharpness in simpler settings (Bartlett al., 2023; Wen al., 2023; Monzio Compagnoni et al., 2023). However full understanding of why SAM works so well remains elusive (Andriushchenko et al., 2023b; Wen et al., 2024). For example, applying SAM on only the normalization layers (SAM-ON] often improves generalization further despite increasing sharpness (Müller et al., 2024). A plethora of SAM variants have recently been proposed with the purpose of even stronger performance or reducing SAM's computational and memory complexity. We focus on two variants of Adaptive SAM (ASAM) (Kwon et al., 2021) which achieve the overall strongest results in Müller et (2024) (see Appendix F.4 for more details).
Tensor Programs. We build on the Tensor Programs framework (Yang, 2019; Yang and Hu, 2021; Yang and Littwin, 2023; Yang et al., 2022, 2023b), which covers many modern deep learning archi- tectures, optimization algorithms and arbitrary dic-paraameterizaions. Each aeeeparameterizationn is essentially defined by layerwise scaling of initialization variance and learning rate as function of network width. Beyond pure infinite-width limits, the simple VI -scallngg allows depth-scaling in ResNets and unlocks hyperparameter transfer across depths (Hayou et al., 2021; Li et al., 2021; Bordelon al., 2023; Yang et al., 2023b). Noci et al. (2022, 2024) provide infinite width and depth analyses for Transformers with the goal of preventing rank collapse.
Look-LayerSAM (Liu et al., 2022) already considers layerwise perturbation scaling with the goal of preserving good performance under large batch training. However, achieving μPP2 with Look- LayerSAM requires nontrivial layerwise learning rate and perturbation rescaling (see Appendix B).
3 SAM induces vanishing perturbations in wide neural networks
This section shows that under the standard (SAM) update rule, weight perturbations induced by SAM vanish in the infinite- width limit in every layer except the output layer. We later demonstrate that other SAM variants also selectively perturb other subsets of layers. For enhanced readability of some formulae, we use colors to distinguish four regimes of perturbation behaviour: Unstable, vanishing, nontrivial and effective perturbations.
While our theory covers any stable parameterization including He and LeCun initializations, for concreteness and for the clarity of exposition, we first present results for MLPs under pPP:
initialize
W ~ (O, 1/d,n), W E R (0, 1/m) for I E [2, L], WL- (00,
with layerwise SGD learning rates ?i T/t,, YR 7), for E [2, LJ, M/+11 Tjn -1
By analyzing the infinite-width behaviour of the SAM update rule, we show that the training dynamics under standard (S) become unstable as the network width increases. This result is first stated informally below in Proposition and then more formally in the next section.
Proposition 1 (Instability of standard SAM parameterization in wide neural networks). Under JuP with the standard (SAM) update rule and default perturbation given in (SAM), the output function becomes unbounded after the first update step in the infinite-widi limit for any fixed, positive learning rate 1j > 0 and perturbation radius p > O.
Hence, to achieve stable optimization, it necessary to introduce some width-dependent perturbation scaling pn -d for some suitable d > 0. To understand the layerwise scaling behaviour of SAM under this scaling, define the notion of vanishing perturbations.
Vanishing perturbations. The weight perturbation e' perturbs the (-th layer's activations as
a' Jz' d((W' + -
where Jz denotes the perturbation of the I-th layer's activations accumulated from the weight perturbations (e" Jeej in all previous layers. We say a layer has vanishing perturbations dur 0 as the width approaches infinity. This occurs if the weight perturbations in all previous layers are too small when measured in spectral norm, that is if ||e" [//[/=/| ||- + 0 for all L' e [[].
3

---

Inpot layer
Hidden layer
Ouipul layer
W/e/WW
SAN
SAM-SCINLL
30D
256
1024
4195
16354
64
256
1024
4056
16384
54
250
1024
4036
16354
0:23
SAM SGD
0.1
1024
16354
1024 16384
1024
16354
Width
Figure 2: ((SAM) effectively only perturbs the last layer) Layerwise weight perturbations (top) and normalized activation updates | ||2 (bottom) for SAM, last-layer SAM and SGD as baseline across widths after training 3-layer MLP in μP with global perturbation scaling p n n-1/2 for 20 epochs on CIFARIO. 10. Average and CI are computed from 4 independent runs, Perturbations are normalized by the weight spectral norm to measure their effect on the layer's output. Activation updates are normalized by / dim (A*)) to measure coordinatewise updates. We provide more neural network statistics in Appendix H.I.
Informally, Proposition 2 below shows that for every choice of a decay parameter > O, either the training dynamics of SAM are unstable all the hidden layers of the network have vanishing perturbations the limit. The formal results are stated in the next section.
Proposition (Global perturbation scaling unstable or induces vanishing perturbations). Fix p > and t E N. Let ft denote the infinite-width limit of the output function after training an MLP of width n with the SAM update rule (SAM) with perturbation radius pm -n for t steps. If d < 1/2, then output perturbations blow up, and f. is unstable. If d > 1/2, then the perturbations in all layers vanish and ft corresponds the limit after t steps of SGD. If 1/2, then only the last layer is effectively perturbed, all other layers have vanishing perturbations.
Figure 2 shows statistics of MLP trained with (SAM) with global width-dependent scaling pm pn-1/2 versus the same MLP trained with SAM where only the last-layer weights are perturbed and e = 0 for all E [L]. As predicted by Proposition 2, both training algorithms produce equivalent training dynamics, already at moderate width, and last-layer perturbations are scaled correctly.
Remark (Practical implications). According to Proposition 2, for sufficiently wide models, any performance gains from standard SAM are primarily due to applying the SAM update rule to the last layer even with a properly tuned perturbation radius. This implics that, when applying the standard SAM update rule (SAM), one can remove the inner backward pass beyond the last layer and nearly recover the computational cost of SGD. However, may be undesirable for optimal generalization only the last layer is perturbed (Figure 1).
Layerwise perturbation scaling. the next section, we show that correcting the (SAM) update rule to achieve effective perturbations in every single layer requires introducing additional hyperparam- eters- - layerwise width midth-dependent scaling of the perturbation radius. This is similar in spirit to jP which corrects standard parameterization by introducing layerwise scaling of the learning rates. We postulate that achieving width-indeppnddett scaling of both updates and perturbations is necessary condition for hyperparameter transfer under SAM. We also lay all theoretical foundations and derive the stable parameterization, we call maximal update and perturbation parameterization (nP?2 that achieves both feature learning and effective perturbations in all layers in the infinite-width limit.
Figure shows that μp indeed achieves hyperparameter transfer in the optimal joint choice of (n- p). while also achieving the best generalization performance (Table 2).
General perturbation scaling condition. For intuitive understanding and a generalization to other perturbation rules, simple condition for achieving effective perturbations in any layer follows from results: in every layer, perturbations should scale like updates in μP.

---

The reason is that both updates and perturbations are gradient-based V Vw:C= = V (z' and thus low-rank and correlated with the incoming activations -¹- Therefore updates and perturbations introduce the same LLN-like scaling factors, and require the same layerwise scaling corrections. Like Yang et al. (2023a), we can rephrase this condition in terms of weight spectral norms to: For a weight matrix W! E ooaan_outtteeeeeeeee oftaiannnnn its update SW; and its perturbation et, it should hold at all times that
||e/ll- = θ (l1)) = = (IWWlll) - e (v f/ffa_oot/faal
with big-O notation that only tracks dependence on network width (Definition C.1). We discuss the spectral perspective in more detail Appendix F.7.
4 Sharpness Aware Minimization in the infinite-width limit
4.1 Characterization of layerwise perturbation scaling: Unstable, vanishing, nontrivial and effective perturbations
To systematically and rigorously understand the width-scaling behaviour of neural networks trained under the SAM update rule, we propose new class of parameterizations, which we refer to as bed- parameterizations. Motivated by the analysis in Section 3, the class of bad-parameterizations naturally extends parneeerramttiiiiioonn (Yang and Hu, 2021) by including layerwise scaling of the perturbation radius. By setting all weight multiplier exponents aj = 0, we do not need to modify the MLP architec- ture and recover representatives of each abc- parameterizzaionn that capture their essence and condense all equations: Ignoring numerical considerations (Blake et al., 2024), each abc- ade--arameteriizttio is essentially layerwise initialization and learning rate scaling. The effects of weight multipliers
SAM are more nuanced than for SGD or Adam (see Remark 12 and Appendix F.6).
To study the infinite-width behaviour of networks trained with SAM in any ped-parameterizatioon we utilize the theoretical framework of NE@OR T programs (Yang et al., 2023b). We write the two forward and backward passes for each SAM update pestunt/perturbatioon step then descent/update step) using the NE@OR computation rules and rigorously track all relevant scalings as provided by the NE@OR T master theorem. All proofs are provided in Appendix E. The full formal result statements can be found in Appendix D. Further theoretical considerations and generalizations around perturbation scaling are provided in Appendix F.
Assumptions. For clarity of exposition, present main results for MLPs. Their extension to other architectures is discussed in Appendix F.5. For all of the results in this section, we assume that the used activation function is either tanh a-gelu for > 0 sufficiently small. For small enough a > 0, d-gelu (Definition C.9) approximates ReLU arbitrarily well. We also assume constant training time width oc. We assume batch size 1 for clarity, but our results can be extended without further complications to arbitrary fixed batch size as well as differing fixed batch sizes for the ascent/perturbatiin and the descent/update step, sometimes used for SAM (Foret et al., 2021). Considering small perturbation batch size is practical, as it has been observed to enhance SAM's generalization properties (Andriushchenko and Flammarion, 2022).
Definition paaaaeemmmtriiiion).... A brd-pparametrization (e U (cthEeL+) U {di (e U {d} defines the training of an MLP with SAM in the following way:
(a) Initialize weights iid as Wt. ~ ~O, n-22))..
(b) Train the weights using the SAM update rule with layerwise learning rates,
W 1+1
W m - Vw:L (5,; W. + Ex) vt).
with the scaled perturbation e, via layerwise perturbation radii,
Vt |
with er (o)..... 0f oft)), :f== n Www.CWlEs (c). (LP)
E, = pn
W.l.o.g. we set ||v.|| = 01), which prevents nontrivial width-dependence from the denominator. This imposes the constraints: d2 > 1/2 min(bi+1, c1+1), d; > 1 min(bi+1, CL+1) ) for e
L), and d 2 1/2, with at least one equality required to hold (see Appendix E.1.3). The normalization w//lex removes
degree of freedom from {d; he(r++1 via the equivalence (E @ {ddiheez+ iff there exists C' R such that di d; C for all / [L 1].
5

---

Stability. To ensure that the training dynamics of SAM are well-behaved with scale, we require codeparameterizaiions to satisfy conditions of stability. Perturbed weights W = W + + induce perturbed activations x' + 8z' and perturbed output function f((5) := fw: (5). We call bod- parameterization stable (Definition C.3) if the hidden activations have width- independeetetttt scaling (1) at initialization and during training, and neither the updates nor the perturbations de of the activations or output logits f: - ft blow up at any point in training.
For stating the conditions that characterize the class of stable bcd parameterizations, we define the maximal feature perturbation scaling F of bedd-parameterizztionn
min min(br+1. CL+1) + d min(d: I1(1 + 1)).
Similar to the maximal feature update scaling r from Yang and Hu (2021), f describes how much the last hidden-layer activations are perturbed as a function of width, x²' + ōx (n--). Hidden-layer activation perturbations do not explode with width if and only if F 0. The output perturbations not to blow up if and only if d + dr+1 2 1 and by+1 + 2 1. In particular, this implies that any stable be-parametterzation together with naive perturbation scaling di d for all ( E [L+ 1] is unstable due to blowup in the last layer. We formally state the stability characterization in Theorem D.2. Ideally, we will later require indhh-independent perturbation scaling which is attained iff f O.
Effective SGD dynamics. Within the class of stable parameterizations, there are parameterizations in which perturbations in the output vanish in the infinite-width limit at any point during training. In other words, SAM training dynamics collapses to SGD dynamics with scale. We are mostly interested in the opposing class of parameterizations with non-vanishing perturbations. We characterize this class in Theorem 6 and refer to them as perturbation nontrivial (Definition 5).
Definition 5 (Perturbation nontriviality). We say that stable bed- prrametttttrinnnonnn is per- turbation nontrivial if there exists a training routine, t E No and & E Rd such that ff((E) = Sw. (E) - (E) = ((1). Otherwise, the bod--parametrizationn is perturbation trivial.
Theorem 6 (Perturbation nontriviality characterization). A stable bed- parametriizaiin is pertur- bation nontrivial if and only if d + di+1 = =1 or min(bL+1, + 1.
For the class of stable and perturbation nontrivial pnalpparameterization,, SAM learning is both stable and deviates from SGD dynamics. A natural question to ask here is: what should be the ideal SAM behaviour in the infinite-widut limit? To address this question, we make the following crucial distinction between non-vanishing and effective pertubbatios.
Non-vanishing versus effective perturbations. Recall that the weight perturbation e' perturbs the 1-th layer's activations as
'00' n' + des' = (((W' + + +da +=))(1+++
where dz' denotes the perturbation of the 1-th layer's activations accumulated from the weight perturbations {s" hr")ret in all previous layers. Therefore, perturbations 5x' can stem both from weight perturbations e" in a previous layer l and/or from weight perturbations e in the current layer I. Intuitively, if we perturb layer, we want this to affect the next layer's activations and thereby have a nontrivial effect on the output function. Otherwise one can simply set the layer's s perturbations to 0 by design and not change the learning algorithm in the infinite-width limit. This motivates the definition of effective perturbations, which demands the weight perturbations of the current layer to contribute non-vanishingly. From the weight perspective (+), effective I-th layer perturbations are achieved if and only if weight perturbations scale like the weights in spectral norm, |E'||=/|||W ||+ (1). Without an effective perturbation e' of the I-th layer, this layer does not inherit SAM's inductive bias towards low spectral norm of the Hessian or enhanced sparsity and does not improve generalization performance. We provide empirical evidence for these claims in Appendix H.2. Therefore a distinction between non-vanishing and effective perturbations crucial.
Definition 7 (Non-vanishing perturbations). For 1 E [Z.]. say that stable parameterization has non-vanishing perturbations in the I-th layer if there exists a t N such that dac 2(1). Definition 8 (Effective perturbations). For / E [L 1], we say that a stable parameterization effectively perturbs the 1-th layer if there exists a t E N such that (x) + dz,-1) = (11), where xt + Jz = =

---

In SP /NTP:
d dz.1
In JPP
d + de.
Collapse to SGD
Collapse to SGD
1.0
1.0
uP?
Non-trivial SAM dynamica
Non-trivial
SAM dynamics
Unstable
Unstable
0.5
-0.
0.5
F- min((mett R2-13
saint
124
min d min(
30 02
Figure 3: (Perturbation phase characterization of aed-parameterrzations) Given a choice of layerwise initialization and learning rate scalings {bi, ci)Ee/++1) the maximal feature perturbation scaling F and the last-layer perturbation scaling d + d1++1 completely determine whether bed- parameterization is unstable, has effective SGD dynamics, effective perturbations in some but not all layers or effective perturbations in all layers. In SP or NTP (left), there does not exist choice of perturbation scalings that achieves effective perturbations in all layers, whereas in μPP (right), there a unique choice as provided in Theorem 11.
Theorem provides a characterization of stable bedd-parameterizztions with vanishing perturbations in any given layer.
Theorem 9 (Vanishing perturbation characterization). For any to E [L], the following statements are equivalent:
(a) A stable parametetttiinnnnonnn has vanishing perturbations in layer lo-
(b) A stable bed-paramettizzttion has vanishing perturbations in layer for all 1 < S to-
(c) Ftg == min (binbb CL+1) + d min (dm 11(m # 1)) > O.
It follows from Theorem 9 that any stable bed- pareaaaattttriztioonnn that performs updates in the original gradient direction (i.e., d, = C' for all E [L + 1] for some C' E R) has vanishing perturbations in all input and hidden layers E [L], and the last layer
T. is effectively perturbed if and only if d= 1/2. This covers the case of both standard and maximal update parameterizations with global scaling of the perturbation radius discussed in Section 3. Negating the conditions of Theorem implies that a stable ped-parameterization has non-vanishing perturbations in layer lo if and only if Ti Achieving effective perturbations a stronger requirement for which Theorem 10 provides the necessary and sufficient conditions.
Theorem 10 (Effective perturbation characterization). For E [L.], a stable bed and-parametrization effectively perturbs the i-th layer if and only if min(bb++1:C++l C1+1) + d + di I[[ + 1) = O.
A stable bed- parametrization effectively perturbs the last layer if and only if d dL+1 1.
4.2 Maximal Update and Perturbation Parameterization (uP?)
We postulate that just as the optimal learning rate transfers across widths under jP for SGD and Adam due to vanishinni width-independent feature evolution in all layers, the optimal learning rate and perturbation radius may be jointly transferable across widths additionally the weight perturbations induce width-indeepnddett perturbations of the activations in all layers. Here, we show that, for every stable initialization and learning rate scaling with b 2 1, there exists a unique stable layerwise perturbation scaling which effectively perturbs every single layer. We term this layerwise perturbation scaling {dr he(.t 1] U {d} the Maximal Perturbation Parameterization (MPP). This concludes the phase characterization of perturbation scaling behaviours (Figure 3).
Theorem 11 (Maximal Perturbation Parameterization (MPP)). Consider any stable bed- parametrization (bi]ee(t+) (ahee(t+) (diieetr--1) U {d). If bL+1 < 1, then there does not exist stable choice of (di))e(L+1) (d) that achieves effective perturbations before the last layer. If bi 2 1, then up to the equivalence di dj + C, C R, VI E [L + 1]. the unique stable choice {di [dilte(t41) U {d] that effectively perturbs all layers E [L + 1] is given by
1= 1, 3/2 min(bL41,CL411
1/2 min(bL41,CL441)
d d=11/2.
di
lE [2, L]. | 1.
(I)
3/2
7

---

Maximal Update and Perturbation Parameterization ppP2. To achieve feature learning in ev- ery layer and hyperparameter transfer in the learning rate, μPP is the unique* choice of layerwise initialization variance and learning rate scalings {Dr, (he (Yang and Hu, 2021). Together
with Theorem 11, this shows that there exists a unique? bed-paraamterrization that achieves both feature learning and effective perturbations in all layers, we call maximal update and perturbation parametrization, μP2 for short. Now that we have found parameterization that achieves width- independent scaling of both activation updates and activation perturbations, ju? fulfills essential necessary conditions for hyperparameter transfer occur in both vi and p.
Remark 12 (Achieving μP2 with weight multipliers). Appendix F.6 covers the extension of our results to nontrivial weight multipliers. We show that, for each choice of weight multipliers (ai)e(L++)) there is a unique" choice of perthheprprrmeeeerss that achieves effective perturbations in all layers. But unlike for SGD or Adam, these parameterizations lead to slightly different training algorithms, because differing subsets of layers contribute non-vanishingly to the joint gradient normalization term V CIr r in (SAM) The term || V L couples all layers so that there do not exist layerwise but only layer-coupled equivalence classes for (SAM). Most importantly, instead of adapting (SAM), we can adapt the architecture with the weight multipliers n at W' with
a = 112 I(I = 1) + 1/2 I[[ - L+ 1)
(a-uP22
to achieve effective perturbations in all layers with naive perturbation and learning rate scaling such that all layers contribute non-vanishingly to the joint gradient norm (Appendix F.6). One downside of (a- (Q- that also applies to naive weight multipliers a; 0, is its incompatibility with unit scaling considerations for low precision training (Blake et al., 2024).
Alternative perturbation scaling definitions. Scaling equivalent to (a-pP22 can be achieved without multipliers by scaling the numerator and denominator terms in (LP) independently, and choosing to scale all denominator terms to be width-independdnt (see perturbation rule (DP) and Appendix F.7 for more details). The ablations in Appendix H.4 suggest that this has negligible effect on the optimal generalization performance of juP2 but can be more stable given suboptimal hyperparameters. Gradient normalization in each layer separately is uncommon and performs slightly worse (Appendix H.5). Appendix F.2 discusses further considerations that led to Definition 4.
Trivial, lazy, and feature learning regimes. A small last-layer initialization variance bL+1 2 is required for stable feature learning. Theorem 11 shows that 1r.+1 > is also required for effective hidden-layer perturbations. Beyond this condition, the choice of {br } and {er} is decoupled from that of perturbation scalings (di) U {d) for stable poneraaramtteizaiinnn, because the scale of the activations of a layer is entirely determined by the scale of initialization b, and learning rates q given stability. Consequently, whether parameterization feature learning regime is independent of the choice of di 's provided that all stability constraints are met. A complete characterization of these regimes for the class of le-parameterizztions has been provided in Yang and Hu (2021) and remains unchanged for the class of stable bed- parameterizations. For completeness, formal definitions and the corresponding results are stated in Appendices C' and D.
trivial, in the lazy regime, or in the
4.3 Generalizations to other architectures and SAM variants
Generalization to other architectures. Our results can be extended to other common layer types, that are representable as a NEλOR program, (Appendix F.5). All considered layer types behave like input, hidden or output layers. Most importantly, normalization layer weights and biases scale like input layer weights to the input 1.
Trogrram including all ResNet and Transformer components
Generalization to other SAM variants. We would like find the correct layerwise perturbation scaling without writing out the NEθOR T proorrm for every perturbation rule individually. Formally justified by our proof in Appendix E, we rephrase our equivalent spectral scaling condition (*) from Section 3 to: maximal stable perturbations are achieved in jPP if and only if e' = O(WW).. This condition holds as soon as weight updates 5W' and perturbations e' are both correlated with the incoming activations f for example if both bradient-based. Table summarizes the application of this condition to two ASAM variants that perform well empirically but cannot be written as a NESOR T program. Additional details are provided in Appendix F.4. We demonstrate that these scalings perform well and transfer hyperparameters in the next section. Note that for hidden layers in
Sttritlyy speaking, unique up to smaller lastr initialization bL+1 2 1.
8

---

Perturbed under global scaling?
For effective perturbations with pp
Input, biases,
Other hidden layers
Output layer
Global
Input- like
Hidden- like
Output-
p
like
norm.
SAM
X
X
/
n//2
n
1/2
3/2
Layer. ASAM
X
/
X
1 n 1/2
1
n-
1
Elem. ASAM
/
/
1
1
SAM-ON
n n//22
Table I: (Layerwise perturbation scaling for effective perturbations in jaP) Without layerwise perturbation scaling (left), each SAM variant perturbs a different subset of layers at large width n + 0o, but we provide the unique layerwise perturbation rescaling jP² (right) that achieves effective perturbations in all layers. This parameterization transfers both the optimal 7j and across widths.
uP22 it holds that e' = !n but W -(n-1/2) entrywise, due to large initialization, showing that it is crucial compare perturbations updates to measure weight scalings in spectral norm.
5 The maximal update and perturbation parameterization μPP achieves hyperparameter transfer and improved generalization
In this section, we provide experimental results showing that jP2 achieves hyperparameter transfer in both 17 and p across architectures, and that μP² also improves generalization over SP and μP with global perturbations - even after multi-epoch training to convergence. We train MLPs and ResNets (He et al., 2016) CIFARIO (Krizhevsky et al., 2009) and Vision Transformers (ViTs) (Dosovitskiy et al., 2021) on ImagenetIK K (Deng et al., 2009). While we directly implement paramttrtttiissssnnns for MLPs and ResNets in PyTorch (Paszke et al., 2019), we the mup- sup-package (Yang al., 2022) as a basis for ViT experiments. Pseudocode and a spectral derivation our yuP impllmmmennaainnnn for ViTs, which is equivalent to (a-HP22, are provided in Appendix F.7. All experimental details are stated in Appendix G and all supplemental experiments can be found in Appendix H.
Comparing candidate parameterizations in MLPs. Figure 1 shows test accuracy as a function of learning rate and perturbation radius for MLPs of varying width. While previous j.P-literature mostly focuses on the more immediate transfer in training error, for SAM it is crucial to consider regularizer, so that optimality in test error
optimality test error as the perturbation radius acts typically coincides with suboptimal training error. In jP without perturbation scaling, the regime of stable perturbation radii shrinks (Figure H.6). In jPP with global perturbation scaling -1/2 the regime of stable o remains invariant under width scaling, but there no significant improvement of SAM beyond SGD, so that the optimal perturbation radius fluctuates within its stable regime due to noise. Only yuP² consistently achieves hyperparameter transfer across widths, and achieves significant improvement over its base optimizer SGD in jP at scale. The full hyperparameter landscapes are provided in Appendix H.3.
p-transfer in ViTs. Figure shows that the optimal pertur- bation radius transfers for ViT-S/16 with SAM in yp? While Andriushchenko and Flammar- ion (2022, Appendix E.3) observe diminishing benefits of SAM at large widths SP, here the improvements beyond the base optimizer AdamW in jP particularly large.
0.09
with
6 on ImagenetI K trained
Top Lccursy
0.04
0:58
0.52
p-transfer for SAM variants in μP2 Figure 6 shows that training ResNet-18 18 in pP2 achieves hyperparam- eter transfer in p for all considered SAM variants with varying width. μP with global perturbation scaling (μP- global) has width-invarian stability threshold in p and the optimal clearly shifts toward that threshold. It would interesting to see whether this shift contin- ues with larger width and leads to suboptimal perfor- mance of pP-global in wider ResNets. Table 2 shows that all SAM variants perform similarly well in pP?
20-1
Perurbation rddis
Figure 4: (p-transfer in ViTs) Training a ViT with SAM in μP2 on ImageNet1K from scratch for 100 epochs yields p- transfer and large improvements over AdamW in pP (dashed lines).
some slightly outperforming the best-

---

SAM yiPgglohal
SAM juP
SAM-ON 10122
Elem. ASAM pu2
Pertatution racias
Pertabation calias
Pertariation radla
Portabution ridios
Figure 6: (p-transfer of ASAM variants in juP2) Test error as a function of perturbation radius P after 200 epochs of training ResNet-18 18 in ju2 on CIFAR1O with various SAM variants (see subplot title). CI over 2 independent runs. Darker lines correspond to larger width multipliers. Other hyperparameters are tuned base width multiplier 0.5. jPP achieves transfer in p and large improvements over the base optimizer (dashed lines) SGD in jP with momentum and weight decay.
SAM global
SAM pP2
SAM-ON μP²
Elem. ASAM μP2
SP 97. 97.00+0.00 (+0.96)
97.00
(+0.96)
97.29 =0.00 ((..1)6))
97.15 (+1.11)
97. 97.23+0.00 (+0.97)
97.3440.08 1108
97.32+0.05 (+106))
Table 2: (Performance of jaP²) Average test accuracy i ssaadard deviation across 4 runs (+ improvement of SAM over SGD) for ResNet- 18 with width multiplier on CIFAR1O using SGD as a base optimizer. In bold, all parameterizations within a 2o-CI from the best-performing variant SAM-ON in μP².
performing variant SAM-ON in SP. This suggests that for ResNets, even with proper lay- erwise balance, normalization layer perturbations may suffice, and performance differences in SP are primarily caused by varying degrees to which the normalization layers are perturbed.
Without providing an explanation, Müller et al. (2024, Section 5.3) observe that only SAM-ON and elementwise ASAM sufficiently perturb normalization layers in SP, Ta-
102
SAM JuP
1006 rox scerncy
SGD μP SAM SP SGD SP
ble 1 (left) explains these observations by showing that only these two SAM variants effectively perturb normal- ization layers under global perturbation scaling. Table (right) also provides full control over which layers to per- turb. For transferring the optimal o with SAM-ON in jP, our theory predicts the global scaling P 122)/2 which confirmed by our empirical observations (Figure 6). However, properly understanding the role of normaliza- tion layer perturbations remains an important question for future work. Note that we report results after fine-tuning all hyperparameters. The performance gain of μu? over SP and ulP-global is likely much higher in larger models, for which fine-tuning is infeasible and the lack of feature learning and effective perturbations is more pronounced. Even under optimal HPs, uP² appears to stabilize SAM's training dynamics compared to SP (Figure 5).
1nt
115
100
105
150
175
200
Epochs
Traming
Figure 5: (Stable training dynamies) SAM in jaP² stabilizes training dynamics for ResNet-18 18 with width multiplier 2.
Future work
This study may serve as an inspiration of how scaling theory can be used to understand and improve training procedures in minimax optimization and beyond. To reach a fully practical theory of deep learning, it will be necessary to take data distributions and training dynamics into account in more detail than it is possible with current Tensor Program theory (Everett et al., 2024). Existing Tensor Program theory assumes constant batch size and training time, and does not make statements about generalization. For example, observe that MLPs and ResNets in SP can sometimes display HP transfer in y and p after multi-epoch training to convergence (Appendix H.3.2). This goes beyond the observations by Everett et al. (2024) rates or weight multipliers. This transfer strongly contradicts the infinite-width theory from Yang and Hu (2021) which predicts output blowup under large learning rates, and it shows that the exact conditions which enable hyperparameter transfer in practice are not fully understood. It also remains unclear how to optimally adapt (SAM) when increasing network depth. We plan to address some of these questions in upcoming work.
observe transfer even without tuning layerwise learning
10

---

References
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware mini- mization. In International Conference Machine Learning (ICML), 2022. Cited page 5, 9, 17, 37, 38, 48.
Maksym Andriushchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion. Sharpness- aware minimization leads to low-rank features. arXiv:23005.1629, 2023a. Cited on page 38, 53.
Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flam- marion. A modern look at the relationship between sharpness and generalization. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning (ICML), 2023b. Cited on page 3. 17.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning (ICML). 2022. Cited on page 17.
Peter L. Bartlett, Philip M. Long, and Olivier Bousquet. The dynamics of sharpness-awan minimiza- tion: Bouncing across ravines and drifting towards wide minima. Journal of Machine Learning Research (JMLR), 24(316):1-36, 2023. Cited on page 3, 17.
Jeremy Bernstein, Arash Vahdat, Yisong Yue, and Ming- Yu Liu. On the distance between two neural networks and the stability of learning. Advances in Neural Information Processing Systems (NeurIPS), 33, 2020. Cited on page 17.
Tamay Besiroglu, Ege Erdil, Matthew Barett, and Josh You. Chinchilla scaling: A replication attempt. arXiv:2404. 10102, 2024. Cited on page 17.
Charlie Blake, Constantin Eichenberg, Josef Dean, Lukas Balles, Luke Y Prince, Björn Deiseroth, Andres Felipe Cruz-Salinas, Carlo Luschi, Samuel Weinbach, and Douglas Orr. u-uP: The unit-scaled maximal update parametrization. 1rXiv:2407.17465, 2024. Cited on page 5, 8, 45, 49.
Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. arXiv:2309. 16620, 2023. Cited page 3, 16.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems (NeurIPS), 13:1877-1901, 2020. Cited on page 60.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training strong data augmentations. arXiv:2106.015548 2021. Cited on page 1, 17.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. International Conference on Learning Representations (ICLR), 2020. Cited on page 17.
Yan Dai, Kwangjun Ahn, and Suvrit Sra. The crucial role of normalization in sharpness-aware minimization. Advances in Neural Information Processing Systems (NeurIPS), 36, 2024. Cited on page 17, 39.
Yann N Dauphin, Atish Agarwala, and Hossein Mobahi. Neglected hessian component explains mysteries in sharpness regularization. arXiv:240). 10809, 2024. Cited on page 17.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2009. Cited on page 71,
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. Proceedings of the 34th International Conference on Machine Learning (ICMI., 2017. Cited page 17.
11

---

Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn. Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x 16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICI.R), 2021, Cited on page
Katie E Everett, Lechao Xiao, Mitchell Wortsman, Alexander A Alemi, Roman Novak, Peter J Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Lestie Pack Kaelbling, Jaehoon Lee, and Jeffrey Pennington. Scaling exponents across parameterizations and optimizers. In Proceedings of the 41st International Conference on Machine Learning (ICML), 2024. Cited on page 10, 49, 61.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness- aware minimization for efficiently improving generalization, In International Conference on Learning Representations (ICLR), 2021. Cited on page 1, 5, 17, 37, 48.
Boris Hanin and David Rolnick. How to start training: The effect of initialization and architecture. Advances in Neural Information Processing Systems (NeurIPS), 31, 2018. Cited on page 16.
Soufiane Hayou and Greg Yang. Width and depth limits commute in residual networks. In Interna- fional Conference on Machine Learning (ICML), 2023. Cited on page 17.
Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Rousseau. Stable resnet. In International Conference Artificial Intelligence and Statistics, 2021. Cited on page 3, 16.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human- level performance on imagenet classification. In IEEE international conference on computer vision (ICCV), 2015. Cited on page 2, 51.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016. Cited on page
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Comput., 9(1):1-42, 1997. Cited on page 17.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optima large language models. arXiv:2203.155556 2022. Cited on page 17.
Satoki Ishikawa and Ryo Karakida. On the parameterization of second-order optimization effective towards the infinite width. arXiv preprint arXiv:2312. 2023. Cited page 17.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and gener- alization in neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2018. Cited on page 16.
Yiding Jiang* Behnam Neyshabur*, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where find them. In International Conference on Learning Representations (ICLR), 2020. Cited on page 17.
Jean Kaddour, Linging Liu, Ricardo Silva, and Matt J Kusner. When do flat minima optimizers work? Advances in Neural Information Processing Systems (NeurIPS), 35, 2022. Cited on page 1, 17.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alee Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv:2001.08361, 2020. Cited on page 2, 17.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009. Cited on page 9, 51, 71.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning (ICML), 2021. Cited page 1, 3, 17, 37, 39, 48.
12

---

Yann LeCun, Léon Bottou, Genevieve B OIT, and Klaus-Rober Müller. Efficient backprop. In Neural networks: Tricks of the trade. Springer, 2002. Cited on page 2.
Mufan Li, Mihai Nica, and Dan Roy. The future is gagussann::: Resnets and their infinite-depth-anm and- width limit at initialization. In Advances in Neural Information Processing Systems (NeurIPS). 2021. Cited on page 3, 16.
Yong Liu, Siqi Mai. Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness- aware minimization. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition (CVPR), 2022. Cited on page 3, 17.
Philip M. Long and Peter L. Bartlett. Sharpness-awan aware minimization and the edge of stability. arXiv:2309. 12488, 2023. Cited on page 17.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 11663))77655-7767 2018.
Cited on page 16.
Enea Monzio Compagnoni, Luca Biggio, Antonio Orvieto, Frank Norbert Proske, Hans Kersting, and Aurelien Lucchi. An SDE for modeling SAM: Theory and insights. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning (ICML), 2023. Cited on page 3, 17, 39.
Maximilian Müller, Tiffany Vlaar, David Rolnick, and Matthias Hein. Normalization layers are all that sharpness- aware minimization needs. Advances in Neural Information Processing Systems (NeurIPS), 36, 2024. Cited on page 3, 10, 17, 35, 37, 39, 41, 48, 51, 52.
Radford M. Neal. Priors for Infinite Networks. Springer New York, 1996. Cited on page 16.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems (NeurIPS), 35, 2022. Cited on page 3, 16.
Lorenzo Noci, Chuning Li, Mufan Li, Bobby He, Thomas Hofmann, Chris J Maddison, and Dan Roy. The shaped transformer: Attention models in the infinite depth-and- limit. Advances in Neural Information Processing Systems (NeurIPS), 36, 2024. Cited on page 3, 16.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetoto, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS), 2019. Cited on page 9, 71.
Ben Poole. Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity deep neural networks through transient chaos. Advances Neural Information Processing Systems (NeurIPS), 29, 2016. Cited on page 16.
David Samuel. (Adaptive) SAM Optimizer (PyTorch). https: htttps///giihhub..cc//aaaaa//// Cited on page 37, 48, 49, 71.
2022.
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. arXiv: 1.0/232, 2016. Cited on page 16.
Sungbin Shin, Dongyeop Lee, Maksym Andriushchenko, and Namhoon Lee. The effects of overparameterizattoo on sharpness-aware minimization: An empirical and theoretical analysis. arXiv:2311.17539, 2023. Cited on page 17.
Leena Chennuru Vankadara, Jin Xu, Moritz Haas, and Volkan Cevher. On feature learning in structured state space models. In The Tiihtt- Annual Conference on Neural Information Processing Systems (NeurIPS), 2024. Cited on page 17.
13

---

Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cengiz Pehlevan. Feature-learning networks are consistent across widths at realistic scales, Advances in Neural Information Processing Systems (NeurIPS), 36, 2024. Cited on page 2, 16, 60.
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How sharpness-aware minimization minimizes sharpness? In The Eleventh International Conference on Learning Representations (ICLR), 2023. Cited on page 3, 17.
Kaiyue Wen, Zhiyuan Li, and Tengyu Ma. Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization. Advances in Neural Information Processing Systems (NeurIPS), 36, 2024. Cited on page 3, 17.
Jonathan Wenger, Felix Dangel, and Agustinus Kristiadi. On the disconnect between theory and practice of overparametrized neural networks. arXiv:2310.00137, 2023. Cited on page 16.
Bichen Wu, Chenfeng Xu, Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Zhicheng Yan, Masayoshi Tomizuka, Joseph Gonzalez, Kurt Keutzer, and Peter Vajda. Visual transformers: Token-based image representation and processing for computer vision. arXiv:2006.03677, 2020. Cited on page 51.
Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and generaliza- tion in deep neural networks. In International Conference Machine Learning (ICML), 2020. Cited on page 16.
Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. Advances in Neural Information Processing Systems (NeurIPS), 32, 2019, Cited on page 3, 16, 42.
Greg Yang. Tensor programs ii:: Neural matrix laws. arXiv:2009. 10685, 2021. Cited on page 24.
Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks In International Conference on Machine Learning (ICML), 2021, Cited on page 2, 5, 6, 8, 10, 16, 18, 20, 21, 29, 31, 32, 33. 41.
Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization in the infinite-width limit. arXiv:2308.01814. 2023. Cited on page 2, 3, 16, 17, 19, 41.
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXivv2203.03466 2022. Cited on page 2, 3, 9, 16, 19, 43, 51, 57, 60, 61, 71.
Greg Yang, James B. Simon, and Jeremy Bernstein. spectral condition for feature learning. arXiv:2310. 17813, 2023a. Cited on page 5, 17, 39, 47.
Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs vi: Feature learning in infinite-depth neural networks. arXiv:2310.02244, 2023b. Cited on page 3, 16, 17, 43.
Yang You, Jing Li. Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations (ICLR), 2020. Cited on page 17.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022. Cited on page 2, 17.
14

---

Appendices
Appendix Contents.
A Notation
16
Detailed related work
16
Definitions
18
D Extensive main results
19
Proof of main results
23
E.1 Tensor program formulation
E,2 The infinite-width limit
29
E.3 Concluding the proof of all main results
31
E,4
Analytic expression of the features after first SAM update
33
Generalizations and further perturbation scaling considerations
35
F.1
Overview over choices of di and
35
F2
Other ways to introduce layerwise perturbation scaling
37
F3
Extension to SAM without gradient normalization
38
F.4
Extension to Adaptive SAM
39
F55
Representing general architectures and adaptive optimizers Tensor Programs
41
F.6
Influence of width-dependent weight multipliers bod boddpparameterizations
43
F.7
The spectral perspective on μP2
47
Experimental details
H Supplemental experiments
53
H.1 SAM is approximately LL-SAM in μP with global perturbation scaling
53
H.2 Propagating perturbations from the first layer does inherit SAM's benefits
54
H.3 Hyperparameter transfer
57
H.4 Gradient norm contributions have negligible effects on generalization performance
64
H.5 SAM with layerwise gradient normalization
65
H.6 Test error over the course of training
66
15

---

A
Notation
Symbol
Meaning
P
width, learning rate, perturbation radius activation function. loss function, input and label time 2-norm as standard for vectors, Frohenius spectral norm malrices (also called operator norm)
L. Be)
|||| = [[o]|a [[]]| := [|]]
standard for matrices
[W |-
w?
trainable weights al time in layer weight updates time in layer
ef
weight perturbations time in layer
leell
norn the rescaled gradient the perturbation denominator
M ot
preactivations and activations at time in layer
dxl
activation updates time in layer
det
activation perturbations at time f in layer /
8fc
update/perturbation of the output function time derivative loss w.r.l. oulput lunction al lime
X: L' M4)
dW? =;
weight perturhations at time in layer (with o for consistency)
elementwise multiplication
de, = 0
derivative output function w.r.t # E {h',=*) at time t, normalized to @(1) derivative of perturbed output function w.r.t perturbed (K..
d=SAMa
time i, normalized to @(1)
scaling the activation gradients
01, 0t
update and perturbation scaling of N and .
Oy... Ow 0 Z
update and perturbation scaling of W
limit scaling; under stability, all considered scalings Ô E (0, random variable distributed according the limiting distribution for the entries of the TP vector specified by the TP Master Theorem
Table A.I: (Notation) Overview over notation used in the main paper (top) and in the appendix (bottom).
B Detailed related work
Signal propagation. Our work be seen scaling theory with the goal of preventing both vanishing and exploding signals in forward and backward passes, where the analysis SAM requires considering stability of perturbations in each layer well. In this sense, build often restricted to an analysis at initialization (Schoenholz et al., 2016; Poole et al., 2016; Hanin and Rolnick, 2018; Xiao al., 2020). For scaling neural networks to infinite depth, residual connections have been found to be beneficial for stabilizing signal propagation while retaining expressivity. The simple VT -calling allows depth-scaling in ResNets and unlocks hyperparameter transfer (Hayou et al., 2021; Li et al., 2021; Bordelon et al., 2023; Yang et al., 2023b). Noci et al. (2022, 2024) provide infinite width and depth analyses for Transformers with the goal of preventing rank collapse and attaining a limit that has behaviour consistent with that of moderately large networks.
rich literature,
Tensor Programs. After kernel-based approaches to understand infinite- width limits of neural networks (Neal, 1996; Jacot al., 2018) and applications of mean-field theory (Mei al., 2018), the Tensor Program series (Yang, 2019; Yang and Hu, 2021; Yang and Littwin, 2023; Yang et al., 2022, 2023b) marks the first important break through in the theory of large neural networks. The framework covers many modern deep learning architectures, optimization algorithms and arbitrary praaetteetttiii,nnnnn,, where each abc- aie-parameterizztioon is essentially defined by a layerwise scaling of initialization variance and learning rate as a function of network width. Yang and Hu (2021) propose the maximal update parameterization (μPP and show that it is the unique stable parameterization that achieves feature learning in all layers in the limit of infinite width. In this framework, training neural networks with global learning rate 7) > for all layers and with He or LeCun initialization falls under the category of so called standard parameterization (SP). The neural tangent parameterization (NTP), studied in the neural tangent kernel literature, differs but does not achieve feature learning in any layer, and is therefore less useful to describe the behaviour of finite width networks than jP (Wenger al., 2023; Vyas et al., 2024). Yang and Littwin (2023) characterize stable learning with adaptive optimizers at infinite width into feature learning versus (nonlinear) operator regime. SAM
16

---

is not covered by the update rule definition in Yang and Littwin (2023) since the nested application of the gradient w.r.t. the weights is not a coordinatewise optimizer anymore. Yang et al. (2023a) show that yP is equivalent to the spectral scaling conditions on the weights || | = = ((Vn//////// /u/ m- and || || = @ m Hence Bernstein et al. (2020) would have achieved their goal of an optimizer with automatic update scaling, if they had normalized by the spectral instead of the Frobenius norm and multiplied by v /fan_outt outtffn_iinn in each layer. While recent works have considered joint limits of infinite width and depth (Yang et al., 2023b; Hayou and Yang, 2023), the data distribution has not been taken into account in Tensor Program literature. The study of scaling laws of jointly scaling model size, data set size and training time has predominantly been empirical (Kaplan et al., 2020; Zhai et al., 2022; Hoffmann et al.. 2022; Besiroglu et al.. 2024). Developing theory to inform Pareto optimal trade offs in a principled manner constitutes an important direction for future work.
As an example of scaling theory for second order optimization, Ishikawa and Karakida (2023) derive μP for KFAC and Shampoo. This scaling rule differs from μP for SGD. Similarly, Vankadara et al. (2024) show that maximal updates are achieved by another different scaling rule for non-standard architectures like structured state space models.
Sharpness Aware Minimization. Sharpness aware minimization (SAM) (Foret al., 2021) has shown to be extremely effective and robust in improving generalization performance across a wide range of architectures and settings (Chen et al., 2021; Kaddour et al., 2022). SAM was motivated as an inductive bias towards flatter minima and it has been understood to have an gradient-nom adaptive edge of stability at which drifts towards minima with smaller spectral norm of the Hessian (Long and Bartlett, 2023; Bartlett et al., 2023). However a full understanding of why SAM works so well remains elusive. While correlations between flatness and generalization have been observed in some settings (Hochreiter and Schmidhuber, 1997; Jiang* et al., 2020), other studies have questioned the usefulness of sharpness as a measure for generalization, especially for modern architectures (Dinh et al., 2017: Andriushchenko et al., 2023b; Wen et al., 2024). Applying SAM on only the normalization layers often even improves generalization in vision tasks depsite increasing sharpness (Müller et al., 2024). Adaptive SAM (ASAM) (Kwon et al., 2021) is variant of SAM derived from a sharpness definition that is invariant to weight rescalings with respect to a chosen normalization operator that leave the output function invariant. The results in Müller et al. (2024) suggest that two of the most promising normalization operators are elementwise normalization T (x) = |W*| O x and layerwise normalization To, (z) = |W'|| F F.z. æ. We state the resulting update rules and a scaling analysis in Appendix F.4. A variant of SAM that is often studied theoretically because of its simplicity does not normalize the gradient of the perturbation. Our theory covers this variant too (Appendix F.3), but Dai et al. (2024) argue that normalizing the gradients for the perturbation is crucial. Monzio Compagnoni et al. (2023) find that unnormalized SAM gets stuck around saddles while SAM slowly escapes through additional fiessian-induce noise. This suggests that the additional effort of analysing the original SAM update rule with gradient normalization is necessary for practically useful theory. Dauphin et al. (2024) draw connections between SAM and other second order optimizers like gradient penalties and weight noise. They show that SAM is able to effectively use second order information implicitly using ReLU, whereas the other two methods close the gap to SAM when using GeLU since they require the localized second order information that GeLU provides in contrast to ReLU. Wen et al. (2023) show that worst-case, ascent and average case sharpness are biased towards minimizing the maximal eigenvalue, minimal non-zero eigenvalue and trace of the Hessian, respectively. With an prchitecture-agnosti analysis, they show that 1-SAM minimizes the trace of Hessian like average-case sharpness, for small enough 1 and p. Similarly, the theoretical results by Andriushchenko and Flammarion (2022) rely on the assumption that learning rate 1) and perturbation radius p are chosen sufficiently close to O. Arguably, the empirically optimal choice of and p lies outside of this gradient flow-like regime and has qualitatively different properties (see e.g. edge of stability literature (Cohen et al., 2020; Arora et al., 2022)).
Scaling theory for SAM. Shin et al. (2023) suggest that the generalisation improvement by SAM continues to increase with growing overparametrization. This corroborates empirical observations that performance monotonically improves with scale, and understanding the infinite-widh width limit is not only of theoretical interest but entails immediate practical benefits.
Liu et al. (2022) introduce Look-LayerSAM AM with layerwise perturbation scaling for preserving good performance under large batch training for enhanced training parallelization. They use LAMB (You et al., 2020) for layerwise learning rate scaling for large batch training. The update scaling strategy in
17

---

these kinds of algorithms follows
V Wic
Wh41 -V/=nolW!!l
with some Ø R. - R2 and where Vw L may be replaced by Adam's "..+E' In practice, often simple functions like (xx) max(c, min(x; C)) or Φ(x) x are used. The idea is to ensure that the update has the same order of magnitude as the weights. Look-LayerSAM follows an analogous approach for layerwise perturbation scaling. derivation of μP for LAMB could also yield feature learning in all layers in the infinite-width limit well as hyperparameter transfer. It certainly requires layerwise learning rate scaling. the case ((x) = 2, following a heuristic scaling derivation as in Appendix F.4 leads to layerwise learning rate scalings 7h = 7 = (1) and T = @nn-1// for hidden layers/ / E [2, LJ. With a bounded function like ((z) max(e, min (zr C)), the scalings become 7/1 (n22/3), 1|r++1
and (1)) for hidden layers E [2, LJ. We leave
closer investigation of feature learning and hyperparameter transfer with LAMB and Layer- in SP and μP to future work.
C Definitions
In this section, collect all definitions that do appear in the main text. With minor modifications, we adopt all definitions from Yang and Hu (2021). If not stated otherwise, limits are taken with respect width n + 30.
Definition C.1 (Big-O Notation). Given sequence of scalar random variables fcr, E IR |=== we write c ② (n=") if there exist constants A, B such that for almost every instantiation of e = (Cz E R n== for n large enough, An- S |c,| S Bn-a, Given a sequence of random vectors
[ E [2" n==1 we say has coordinates of size θ (n ()) and write x θ (n -a) mean the
scalar random variable sequence { [VIl,,11 /n is θ (n= For the definition of O(n
and c = Q(n "), adapt the above definition of c= O(n ") by replacing An S |cal S Bn m
with |c,l < Bn and An < |cnl, respectively. We write In = o(n ")) if n" /2 2 /n - 0
almost surely.
Definition C.2 (Training routine). A training routine is a combination of base learning rate 1) 2 0. perturbation radius p O, training sequence {(E,, Je) and a continuously differentiable loss function ((f(E),) y) using the SAM update rule with layerwise perturbation scaling (LP).
In addition to the stability conditions from the corresponding SGD result, we demand that the activation perturbations do not blow up. Otherwise the perturbations would strictly dominate both the initialization and the updates which makes the perturbation too strong and is avoided in practice.
Definition C.3 (Stability). We say bod bed-parramttrization of an L-hidden layer MLP is stable if
1. For every nonzero input E (d))
N -0 (1),, VI E [Z]. and BS6(S)²
where the expectation is taken over the random initialization.
2. For any training routine, any time & E N, / E [Z], E Reda, we have
hE(E) =(( =5) = O.),,, and Sτ(E) = O.
where the hidden constant in O. can depend on the training routine, t, E, I and the initial function fo-
3. For any training routine, any time t E No. [L], E E Rd,, for the perturbed (pre-)activation ii; = 1 (Wa), = x! (W, and output function f.(W)) we have M(()) - n1( - =E( = O.(1). and f((E) = (1).
where the hidden constant O4 can depend on the training routine, t, E, I and the initial function fo-
18

---

Definition C.4 (Nontriviality). We say perametretiiizotooonn is trivial if for every training routine, fr(E) - fo(E) - 0 almost surely for n - oc, for every time >> 0 and input & E R Otherwise the wrd-parametrization is nontrivial.
Definition C.5 (Feature learning). We say a bedd-paramerrization admits feature learning in the I-th layer if there exists training routine, a time / > 0 and input such that (5) (5) = Ω. (1), where the constant may depend on the training routine, the time I, the input & and the initial function fo but not on the width n.
Definition C.6 (Vanishing perturbations). Let LE E [L]. We say that stable bed- parametrization has vanishing perturbations in the 1-th layer if for any training routine, t E No and & E Rd,, it holds that F 244 (1), and it has vanishing perturbations in the output if for any training routine, t E No
and { E Rd it holds that fc(E) == fw. (E) - fw,(5) = =[1).
Definition C.7 (Perturbation nontriviality). Let I E [L]. We say that a stable bod- paraamttttiiiooo
perturbation nontrivial with respect to the 1-th layer if and only if it does not have vanishing perturbations in the 1-th layer. A stable bed- padamaaaamtiinniooon is perturbation nontrivial with respect to the output if it does not have vanishing perturbations in the output.
Definition C.8 (Effective perturbations). Let 1 E [L. + 1]. We say that stable bed-paaramttiization effectively perturbs the I-th layer if there exists a training routine, & E N and & E Rd such that OW 11
(E) (E) (11) where SW; is defined in (LP) and z x. E1-
Definition C.9 (o-geluu. Define o-gelu to be the function x + (1 +
eef erf (o=1z)) +
2/7
In order to apply the Tensor Program Master Theorem, all Nonlin and Moment operations in the NE@ORT program, which do not only contain parameters as inputs, are required to be pscudo- Lipschitz in all of their arguments. For training with SGD, this is fulfilled as soon as ø' is pscudo- Lipschitz. Both tanh as well as o-gelu fulfill this assumption.
Definition C.10 (Pseudo-Lipschitz). A function f R* - IR is called pseudo-Lipschinn of degree d if there exists C >>0 0 such that lf(x) [(v) S C|| wll(1 + Σ |r:|d + () We say f is pseudo-Lipschitz if so for any degree
D Extensive main results
Using the formal definitions from Appendix C, here we provide the full formal statements of all of our main theoretical results together with further details and implications. The proof of all statements is provided in Appendix E. Since SAM evaluates the gradients on perturbed weights, it is not covered by the update rule definition in Yang and Littwin (2023) and an infinite-width analysis requires explicitly deriving the corresponding NETORT T program, scalings and width limits.
Recall that our definition of bermttarrttttiiizaiionns extends abc- parameterizaaions by setting the maximal perturbation scaling to n -d and allowing relative downweighting n di of the global scaling in each layer l. The perturbation scaling does not affect the choice of layerwise initialization variance scalings by and the layerwise learning rate scalings q, Common be- parametrizations for SGD are summarized in Table D. SAM with SGD as base optimizer requires the same scalings. Similarly, SAM with Adam as a base optimizer requires the same scalings as Adam (Yang et al., 2022, Table 3). Recall that, for convenience, require width-independent denominator scaling ||u,1| (1) of the scaled gradient for the perturbation (LP), which imposes the constraints
di 2 1/2 min(b1+1: CL+1), de こ min(b C1+1) for I E [2, L), dr+1 2 1/2. (D.1)
All (pre-)activation and function outputs can be thought of as outputs given a fixed input & Rd= with d,, E N fixed, e.g. ft := fw. == fw.(E). For the perturbed weights we write W, := W, + SW, with 5W; defined (L.P) as e/. Here we write weight perturbations as 8w? instead of e/ to show the resemblance weight updates SW/. Perturbed activations and function outputs at time are written as { = - (E) and fe(E) = Swz (E). Recall that for all of the results in this section we make the following smoothness assumption on the activation function.
Assumption 1 (Smooth activation function). The used activation function either tanh or a-eelu for a > 0 sufficiently small.
19

---

We define the maximal feature update scale of bcd parameteriizzioonn
min C2+1, 1 + min(a I(I * 1)).
(D.2)
as well as the maximal feature perturbation scale a bedd-parameterization
i = (iin + d min - II(2 1)).
(D.3)
1=1
Stability requires the constraints (a-c) from SGD and additional perturbation stability constraints (d-e) that include the layerwise perturbation scales {di)t=1....L-1.
Theorem D.2 (Stability characterization). A bed-parramerrization is stable if and only all of the following are true:
(a) (Stability at initialization, h/. = (1) for all I, fo 0(1))
bj 0, = 1/2 for [2, Li and bL+1 2 1/2.
(b) (Features do not blow up during training. ie. Ar, = (1) for all 1)
r 2 O.
(c) (Outpul function does not blow up during training, i.e. AW. w Aut 0(1)
CL+1 2 and b + r 2
(d) (Feature perturbations do not blow up, i.e. dat ((1) for all
Z O.
(e) (Output function perturbations do not blow up during training, i.e. öw.++1 2/, WT+ Ez (()) d+ dL+1 2 and bL+1 + F F 1.
The nontriviality and feature learning characterizations from SGD remain unaltered. This is because in the definition of r, it holds that d d+ddr. dr,+1 2 (from perturbation stability), and min(br+1s C1.+) )) S already had to hold for nontriviality in SGD, so that stable perturbation scaling does not affect Theorem D.3 (Nontriviality characterization). stable bed- sad-parametrization is nontrivial if and only if CL+11 = 1 or min(b CL+1) 1.
As for nontriviality, the conditions under which a stable, nontrivial parameterization is feature learning in the infinite-width limit are decoupled from the choice of perturbation scalings {dhhe(t+] U {d}. Hence the conditions are the same as for SGD, Below provide slightly refined result in terms of the maximal feature update scale Tlo of bod- parameterization up to layer lu (as provided in the Appendix of Yang and Hu (2021)).
Theorem D.4 (Feature learning characterization), For any lo E [L], the following statements are equivalent:
(a) A stable, nontrivial bedl-arrametrizatoon admits feature learning in layer lo-
(b) A stable, nontrivial ped-parametrizzation admits feature learning in layer for all I lu-
(c) Ti, min(bL+11 CL+11 + dL++1 + min (Cm Il(m & 1)) O.
Consequently, stable, nontrivial bed- paraantrtetiiinioonn admits feature learning (at least in the last layer activations) if and only r= O.
Remark D.5 (Effective feature learning). As for perturbations, feature learning in later layers can be caused by weight updates in earlier layers that propagate through the network. One could demand effective feature learning in the (-th layer as ow
(11) and would occur if and only if
min(br+1: C1+11 d d1+1) a I1[ + 1) 0.
As for nontriviality, perturbation nontriviality in the output is attained if the constraints for SwL+1af Wark exactly satisfied.
Theorem D.6 (Perturbation nontriviality characterization). Let E [L]. A stable bed-
parametrization is perturbation nontrivial with respect to the 1-th layer if and only
ft = min(be+1, (1+1) + min(dm. I(m 1))
A stable beal-puramttrization is perturbation nontrivial with respect to the output if and only if d + d = 1 or min(bL-1,, CL+1) 1.
20

---

The converse formulation of the perturbation perturbation-nonttiiialli results characterizes the regime of vanish- ing perturbations.
Corollary D.7 (Vanishing perturbation characterization). For any to E [L], the following state- ments are equivalent:
(a) A stable bed-parann ond-parumetrization has vanishing perturbations in layer lo-
(b) stable bed- paraaetttriiitoonnn has vanishing perturbations in layer for all 1 S lo-
(c) Tio := min(bb++1, (141), + + min ( I1(m + 1)) >
A stable bed-parde parametriization has vanishing perturbations with respect to all layers and the outpur function if and only if d1+1 1/2 and f max(0, bL+1). This case reduces the results in Yang and Hu (2021).
For perturbation nontriviality it suffices that the perturbation in any of the previous layers is scaled correctly. For effective perturbations, we need the correct scaling in exactly that layer.
Theorem D.8 (Effective perturbation characterization). For I E [L]. a stable bed bed parametrizattio effectively performs SAM in the 1-th layer if and only if min(bb--1 CL-11) + d + di I(I + 0.
A stable bed- andaparametriztioon effectively performs SAM in the last layer if and only if d dr,+1 1.
The above understanding of all update and perturbation scalings allows us to extract the most important consequences of different choices of perturbation scaling on the learning dynamics. Beyond vanishing hidden layer perturbations, the following theorem shows that the joint gradient norm | can be approximated efficiently without an additional backward pass under global perturbation scaling.
Theorem D.9 (Global Perturbation Scaling). Given any stable bod- parametrrriiiio {bi (E))
{a heuLt U {di (di U {d). The parametrization performs updates in the original gradient direction if and only di C for all E [L + 1] for some C E R. In this case, the parametrization all hidden layers E [L]. and the last layer I L + is effectively perturbed if and only if =d==1 1/2. lfbL+1 > 1/2 (as in μP), the gradient norm is dominated by the
has vanishing perturbations
last layer and simplifies to,
|| = 1n1112-0)
|p:|| - (E),))))
=0(1
One might suspect that it is desirable to let all layers contribute non-vanishingly to the gradient norm in the denominator of (LP). The following proposition shows that this should be avoided with our definition of bed- bed-barameee parameterizations. Of course, if we add even more hyperparameters by decoupling set all contributions to e(1), which is what we do in
numerator and denominator scalings, Appendix F.7.
Proposition D.10 (Balancing gradient norm contributions). Given any stable craaettramttioniooonn {bi he(t+1] U (ci)Ee(L++) U [d()eett))] U {d]. If all layers contribute to the gradient norm non- vanishingly in the limit, i.e. ||l||| = ([ II) for all E [L+ 1], No. then the parametrization has vanishing perturbations in all hidden layers 1 E [L]. Such parametrization effectively performs SAM in the last layer [ L + if and only if 1/2.
The following theorem provides the unique correct perturbation scaling for any stable be- parameterization with b 2
Theorem D.11 (Perturbation Scaling Choice for Effective Perturbations). Given any stable nodd-paramttization {hr hejL+4) U (cthEEL+1 U {di U (d] If bi < 1, then there does not exist a stable choice of (di)eeu))) U {d) that achieves effective perturbations before the last layer. If b 2 1, then up to the equivalence di = di + C, C E R. VI E [L+ + 1], the unique stable choice [d] U {d} with effective perturbations in all layers I E [L + is given by
1/2 min min(bL+1,,C41)
1 1, E T.] = +
d 1/22,
di
3/2 min min(b141,CL41)
(D.4)
3/2
In this parameterization, the first layer dominates the gradient norm as
|| O(1),
||l| || - ||ell|| (
21

---

Definition
SP
SP (stable)
NTP (stable)
uP
0 1= 1
0
1 1,
0
1= 1,
0
1 1,
N(
1/2 IE [2 Li,
1/2
12 2.
1/2
122 2.
11/2
122.
1, 1= I+
-1 1=1,
s0 1=1,
Ct
LR 77
0
[1 122
0
1E 2, L.
(1
1= L+1 0
Equation (D.2)
-1
1/2
1/2
Stable?
Nontrivial?
Feature learning?
Table D.I: (e-pparametrizations) Overview common implicitly used ce-paraametrizations for training MLPs without biases in standard parametrization (SP), standard parametrization with maximal stable nonadaptive LR (SP (stable)), neural tangent parametrization (NTP) and maximal update parametrization (μP)).
Definition pm
Naive
Global (stable)
Effective
d
0
1/2
112 1/2 cv
1= 1,
di
V WiLL
1/2
1/2
3/2 cp
1 [2, L},
3/2
1= L+1.
F
Equation (D.3)
CV 1/2
0
Stable?
X
Last layer effectively perturbed?
All layers effectively perturbed?
X
X
Table D.2: (Perturbation scalings) Overview over important choices of the global perturbation scaling on -d and the layerwise perturbation scalings n -: for training MLPs without biases with SAM: Naive scaling without width dependence (Naive), maximal stable global scaling along the original gradient direction (Global) and the unique scaling that achieves effective perturbations in all layers (Effective). An extensive overview that characterizes all possible choices of perturbation scaling is provided in Appendix F.1. Recall the gradient scaling cv := min(br-1.. CL+1).
Table D.2 summarizes the consequences of Theorem D.11. Together with Theorem D.11, the following proposition suggests that 6ı+1 is good choice. However b > 1 can also induce effective perturbations, as long as and d are chosen correctly.
Proposition D.12 (Effects of last-layer initialization b on all perturbations). If stable bed- parametrization with min min(b1+1, (1b1+11,C44)) I 1 is perturbation nontrivial with respect to any hidden layer l E [L), it is also perturbation nontrivial with respect to the output.
Lastly, the following proposition shows that effective perturbations from the first layer propagate through the entire network.
Proposition D.13 (Perturbations propagate through the forward pass). All stable bed- parametrizations with d1 = 1,i(br.+1,,t+1)) - dd effectively perturb the first layer and are perturbation nontrivial in all layers.
Remark D.14 (Efficiency gains). The above results may be used for efficiency gains. Given any stable parametrrizzai,nn we can compute the maximal layer lo such that fl, 0, and in wide networks do not have to compute SAM perturbations before layer to + 1; soon as bL+1 1/2 (as for μP), the gradient norm for the SAM update rule is approximately given by IVL & (f+(()) Ie which can directly be computed without an additional backward pass. The practical recommendation from our experiments however to either use juP2 to completely abstain from perturbations.
Remark D.15 (SAM without gradient normalization). For the SAM update rule without gradient normalization simply set d 0 and remove the gradient norm constraints (D.1) to arrive at the
22

---

adapted NESORT program and bed-constraints. Note that standard parametrization gets even more unstable without dividing by | @ now requiring d++1 2 1 for stability, Similar to the previous results, this shows that unawareness of pedepaaramttiiationns requires strongly scaling down p for stability, while vasting computation on vanishing perturbations before the last layer. More details can be found in Appendix F.3.
E Proof of main results
In this section we derive the NERORT program that corresponds to training a MLP without biases with SAM. For simplicity and clarity of the proof, prove the one-dimensional case din 1, dous = 1, but an extension to arbitrary but fixed din, dout is straightforward. Recall Assumption that allows us to apply the Tensor Program Master Theorem and explicitly state the widtn limit of training MLPs with SAM in Appendix E.2.
E.1 Tensor program formulation
E.1.1 Tensor Program initialization
We initialize the matrices W....., W as (W)as ~ (O, which absorbs by 1/2.
We initialize the input layer matrix Wi E Rax1 and normalized output layer matrix W Wo n E 121xn as (W))a. (W (W,++)) ~ ~NO0 1), as initial vectors should have a distribution that is @(1).
In the NE@OR T formulation, we write all quantities 0zz, where O2 denotes their scaling n for some C' E R and = therefore has ((1) distribution. The stability, nontriviality and feature learning conditions then stem from requiring either 0g - 02 depending on and its desired scale.
E.1.2 First forward pass
We denote a definition of Tensor Program (TP) or NE@OR N coomputatio as ==. Compared MLPs trained with SGD nothing changes in the first forward pass,
(E) = 5 (NL); of = ((h)) (NL), n == W 2/- (MatMul)
In the case of MuP, fo({) = W -/ (E) + 0 defines a scalar in the TP.
Observe the scalings =1 (h)) (nn +)) ), =nj = O(h)) 1(2 for E [2, L] due CLT, independence at initialization and w/ O(h)) = ((1) by stability. Hence stability at initialization
inductively requires b, = 0, by = 1/2 for l E [2, L] and br 2 1/2.
E.1.3 First backward pass
The chain rule of the derivative remains the same, we just evaluate on different weights compared to standard SGD. We denote the adversarially perturbed weights by W and the normalized perturbations by 5w! Before computing the updates we have to compute full backward pass to determine these perturbed weights for each layer, and then compute forward pass with these perturbed weights to compute the perturbed preactivations iς that we will need for computing the SAM update. Therefore the NE@OR program for SAM maintains a perturbed copy of all preactivations, activations, last- layer weights and logits just for computing the updates of the actual parameters.
Under MuP, the loss derivative with respect to the function remains Xo := L' ('(ulEE)) Mo) - Xo = L' go). For the weight perturbation, need perform SGD backward pass,
dxfp Wo!+. dhl = day ø' (((h), dap = (W))Tank,
where dz = V''V... For SGD (and for SAM, as we will see later) all gradients have scaling O- n -b2+1 in the first step, whereas overload the notation 0r later steps. For clarity presentation assume br+1 2 CL+1 here, the other case follows analogously. For the first step this can be understood from
min(b2+11.244) for all
Vz fo W (n bLt1),
WA- Vzr o ('( O(n -bat),,
23

---

since h/ = (1) by the stability assumption, and this scale 111=b111) propagates through all layers via the chain rule and remains stable in later backward passes. For hidden layer gradients, observe that
V Val 1f
= (W + AW
1-1
+ (W))) Vafee - n ((Vnnfo) fa) V fe)a
==
O 1n1222b
2b 0: n (n)) (
where first term's scale stems from the products (w 0
v) due to Yang (2021),
dy 1/2 for stability at initialization and 11+1 cL 2 1 for update stability during training (r 2 0)- we allowed the second term to strictly dominate, the gradient scale would explode iteratively in the backward pass.
The gradient norm. Before computing the weight perturbations, we need to compute the gradient norm for the SAM update. The gradient norm at time t in each layer 1 E [2, L] is given by the scalar,
0L awi
8-²
= Σ (x)(d))) = (x ( 1 = x² ((dh))) dh:) ((o)-))aat G-=1
where = L' ('()(EE), m1) and we used 0h /ow!, = (1r) S5x ( I=== I....
Hence the gradient norm of all weights jointly is given by the unnormalized scalar
V 1 = xi (no2 (dh))Taht dhd (5T
(dh))² dhl (x)-")Tart
1-1
( Itf
+ Σ n 02
91 +
n
n
n
2=2
(E.1)
with scaling |i = Onn2f + n2) = (n)) because stability at initialization requires P 2 1/2 so that n²06 S n. Note that the first layer contributes vanishingly to the gradient norm, the hidden layer gradients only b = 1/2 (equivalently fo @()) and the last-layer activations always in dominating order. So in jaP, in the limit, V L' (f)(()))) v.) Ia This means that the unscaled gradient always aligns with the last-layer activation. For learning in jP, this dominance is corrected by the layerwise learning rates.
The squared norm of the rescaled gradient is given by
n0 2dh (dh!)Taht dh
||0 xi
(ET Ev)
(E.2)
n
L
2d, (dh)) (xt (x!-1)Ta,t 1-1 n R
xt
20 11 (x.
nn
n
n
n1
1=2
with scaling 83 = -n 1- 2d 07
+ ン n² 22-2d 02 + n 2d1+11 ). For simplicity, set 0, = 1. This
raises the constraints n2 221 0z rewritten as
S 1, n
2- 2d
oZ I 1 for E [2, L] and n 1d22 < 1, which can be
di 2 1/2 min())+1,
d; 2 1 (01-min(b4+i,c444) for / E [2, L),
dr.11 2 1/2,
where at least one equality is demanded to hold in order to attain 0, = 1. If one of the equalities holds, the respective layer contributes to the norm non-vanishingly in the limit.
Thus, applying the square root and dividing by 0, the square root of (E.2) defines normalized TP scalar.
Stability implies that also the perturbed (pre-)activationn and output function remain ((1) and O(1) respectively. Otherwise a SAM training step would induce blowup in the updates. We call this weaker property of just the perturbations perturbation stability.
Perturbations.
Definition E.1 (Perturbation stability). We call a paaaterraetiinoiioonn perturbation slable if and only if it o (1) for all e [L] and e N and of ((1) for all t N.
24

---

Mathematically we get the normalized weight perturbations for 1 E (2....,L},
5Wo+1 Xo o6
Xo dhg (xb ITT
Sw Xo dhi co
5w
||vol|
000
|100||
which scale as Oc+1 Owars := n (2+dx+1), ((n(dta)-brt)) and ((1=))4)))-b respec-
tively. But the NEgORT program computation rules do not allow to compute matrices 8W E [L], therefore we use the weight updates to directly compute the preactivation and activation changes analogous the f-th forward pass. For all f. 2 O, we write
hi h4 + 0,ō with the perturbations for E [2, Lj.
zt = xt 0:804,
PXo E)
(h)(E)
|
824 :=
8, (p(), + 0,8h,) ((h))),
0,5 :
Wi + (W: - W
0 W 1-1 pāw; Xo 12
dhg.
|voll n
which defines NonLin operation with the vectors W/ 1-1 and dhg and everything else treated as scalars, and with first backward pass sealings Ow == n -n-nnddde (d++1) Ow: == n' n 0r and 0. == max(@ 0w:) = max 6w,, where we used that s01 ((1) due to perturbation stability. Note that these scalings may implicitly increase when replaced by 0g min(b.+11.t.+:
since Ax n bb-b gets
The activation perturbations can then simply be defined via the NonLin operation,
Bano 0, (0(h) 0,5%)) <(b()),
with the same scaling 566-
The perturbation of the scalar output function can simply be defined via the NonLin operation,
öw ah
W 82g
8 Wo+1 2 Wh 2t =0
+ +
n
n
with 01+1 == n0wz+ and 8/ == n0
SAM Update. Finally, compute the SAM updates follows. In the case min(0L+1,,C411) s + d the weight perturbation scale is dominated by the weight scale, so that
d = Wi + U Sw!++
with (141// = 1 < 1, whereas if miin d + dr+t we write
dæsAMO W W/6 öW
with (0/(+)))) := 00/ 0-/0 S 1. In any case, the scaling of da AAMO and all other SAM gradients 0saM = (0 (d+d2+1)) min(b241.c144 d+d1+1), The other SAM gradients given by
dh'AMM.0 :
dr's SAMoO @ ((hh)) (W)) dh's = (W! + ōw:)) dh's Xu (dhb)" dh'sAMM 2
dr SAM.0 =
(W!) dhh' + poSAMs
|!ool|
where the last line define NonLin operation in the vectors (W:) Th SAM, and 2// and everything else treated as scalars. Consequently, V flw, is of the same scale as V vag J'w,, and V Vlw,fw, is of the scale max(0sAM @w:0sAM) 0saM since 0 < 1 is required for perturbation stability.
25

---

Note that for SAM's weight updates the loss derivative is also evaluated on the perturbed weights,
Ko = C' : yo)-
Constraints on the output function. Assuming z/ = (1) (perturbation stability), we get Xo = O(1) if and only if 5W = -n if and only if + d1+1 2 1.
We have Xo = (1)) if and only if Jo W (1). This can either be caused by changes in the last-layer weights, by non- vanissnisii initial function W x (if and only 7 = 1/2) or by W 82f (1), which holds and only if b + Fr. = 1 (analogously, W oxg = 0(1) if and only if bL+1 + FL 2 1). The first case requires äw?++
O(n "(, since oW and 2t O(n ((nn") ') is fulfilled if and only if d + d#7b (the analogue to CL+1 2 1 for stability and CL+1 = for nontriviality).
highly correlated. OW
Hence perturbation stability of the output function holds only if d + d 2 1 and 1 + + 21 1. Then, perturbation nontriviality holds and only if d d=+1
bL+1 + Fg = 1.
In the 1-th backward pass, bu+1 + Fz 2 1 will be replaced by the slightly stronger constraint
bL+1 + - 1.
E.1.4
1-th forward pass
Formally, sum the updates in each step,
Wr++1 : -W + OLTL/ +....884+t))
where SW (41 = -1 Xt (zt)² denotes the normalized change in the weights WL-1 (as a row vector) of scaling bz+1 = Owr+: = n under perturbation stability and nontriviality so that W scales 0r
min(be11:<2 SWI++ should not be confused with 8w/4+ which r+1 denotes the perturbation of the weights at time + For every nontrivial stable parametrization we have Xt (1) and o = @(1) which requires 0g S 1. In the case CL+1 < b we write W : n bu +er+1 Wo ++ + (6W7++1 ..... &W?++l) with the same scaling By = n (nin
For preactivations and activations we also sum the changes from each step,
M = No + Oc(h + oh1), xt== o' + Or( + +... + +zx)).
Using the fact that
W - W - 10 -0
yields the normalized preactivation updates
(E) == -X--Idh (NL),
with scaling 01 Oy1 n ASAM no - - as for SGD under perturbation
stability and nontriviality where Xi-1 = (1).
For 1 E [2, L). it holds that
W W/ W1 1xe- 1 dh
- AMM.tA -
with the right scaling Ow = n² 1-----minn oi
as for SGD under perturbation stability
=/=1 = @()), so that we get oh, using a telescope sum,
0:5h W! W of
Whi x + (W! W
(x)
11 Σ (w! W 12-1)or + (W: W
8:-1
Wij00t
1-1
(+ -oot
01-1
110
Xa-i
dh's
n
1=1
26

---

(z(-1)(=+=-
w Xe-1
dh
which defines a NonLin operation with the vectors Woo .dl'sAM.. di'sAM,.-1 and everything else treated scalars. The scaling given by
0 max 01 0w:) ( max Pwm = n
m=1
with
Ti = min(bb--, CL+1, d + d1+1) + min min I(m + 1)),
m=1
where Ow < 1 for all 1 E [T.] for stability. Note that for I I l2, it holds that 01, S 01g which explains the sufficiency of 0, n TL = n for the stability of the activation updates.
Activations with the same scaling 8 can then simply be defined via the NonLin operation
d1t = 0, (o(h)_1 + 0:5h.) o(hk-))).
The updates of the output function scalars defined
8ʃc =- 0/41 SW of
WL++
dat
0/.
1-1
where 0/,11 = n0 = -++1 and // = nl00, = n² m where we will see
why W (n min(br+1.ez=1)) in the next paragraph. This leads to the constraints Cr++1 2 1 and b1+1 + 2 1 for the stability of the output function, where equality in either constraint leads to nontriviality.
E.1.5 1-th backward pass
Perturbations. Due linearity and stability, the last layer remains
dxy == W
with scaling 0x = m114))
As in the first backward pass, use the weight updates to directly compute the preactivation and activation perturbations similar to the /-th forward pass but performing SGD instead of SAM the last step. The SGD backward pass for the perturbation is given by
dhp do--1
dat o o' d'(N), (W!)² anl
Wil nn:w Σ xe-1 n
dh (
dhp
t Xo-1
(dh's :Thh
dh
Wydhh nnnn2 ABAANS
1-
n
with scaling max(Oe (0v, n² (1~=00AM0r) = or, since n 05AM & 1 is implied by 2 0 required
for the stability of (pre-) pare-vactivation updates.
We write X (fe( y1) for the derivative of the loss with respect to the unperturbed function (which is ((1) ) under stability and nontriviality), and get
+ Pxe(E? Sdh ||
0,8h :
0, W680 + +Wi W.)Jt(1
1
0,-1
W6boot- +
Σ (W: - w w + (W! W/)of
s=1
27

---

( but
ii-1 W6bbst-1 n(nn 8SAM) Xs-1
dil's al'SAMio-)
n
(
+pBw:
Xt
dh
||10ll
which defines NonLin operation with the vectors Wo ddh' SAM, dh A dhl
and where we can now define the definitive scalings 01 = Ow == n2 (d+dd 0- =
n 000Mm 0w:) max Bwm =n with
Owr = n -d-(d++ = n
and B1
f = min(br+11, + min(dm 1mm + 1)),
m=l
where we used that n 0SAM
due t 0 for stability and z! ((1) due to perturbation
stability. Perturbation stability of the hidden layer sareevaativations 5h', 8z (()) for all I [4] holds if and only if == Fz. 2 0 since ñ 2 Fr. for all I S L.
The activation perturbations d and the perturbation of the output function ff can be defined exactly as in the first backward pass,
dak =
( + 0,0h)) - -((1)),
SWLtist + +0L0 n
of
Wh+1it- Wh+1 - =61+1 n
with öw := e x o
and the same scalings 0, 0/411 = n (d+ddd- and 07 = ndv0, m1 since W W + AW max(n - bu+z n o2+1), which yields the slightly stronger constraint (than in the first backward pass) min(b (++1) + 2 1 for perturbation stability and either 8/11 1 or min(bL+1,
well
F 1 for perturbation nontriviality.
SAM Update. For each E (I,...,,), as in the first backward pass, we get
detAMM = W + O2+1//V Sw,+1,
with scaling OSAM = n d1 well as
dh's de'sAM, + d' (k))..
For dæ'SAMA SAM, we again use a telescope sum over the weight changes,
dx SAM, =
(W!) dn' SAM, (Wy) Owi Σ ow! @widw!) dh's d
(dh'sAM.o S ddd'sAMa +111
TW dh dh' - innn'=90AAN Xa-1
n
(dht)" dh dh'saMe 1-1
+p(n2/2240 112--007) Xe
Ieelll
which defines NonLin operation in the vectors (W))² dhl SAM, 30-1 erything else treated scalars. Note that the scalings remain OSAM, since / W.
+|-1, "...,,,-_,,11 and
O O(max(PsAM,, - n/2-4d((0sM)) ((fSAM) under stability, nontriviality, perturba- tion stability and perturbation nontriviality.
Finally define the loss derivative on the perturbed output function
:== CWWWL+124 yy),
and compute the normalized change in WL+1
Sw [+1
-
28

---

E.2 The infinite-width limit
In this section, apply the Master Theorem's 's computation rules derive the marginal distributions Z corresponding to the vectors of the program constructed above. According to the Master Theorem, each such vector will have roughly iid coordinates distributed like 2a in the large limit.
We assume stability holds, so that 0 - 0 E {0, } for all scalars A in the program.
For the first forward pass, we have
7ho(e) E7W6
zrole) = (Z
zd (5) 7W *b((6)
If b > 1/2 then fu O, otherwise if DL+1 1/2 then fo converges nontrivial Gaussian. For the details we refer to Appendix 1.4.1 in Yang and Hu (2021), as at initialization their results still hold here.
For the first SGD backward pass, we have
zdz/(0) (0) = 2 Wo!
zani(e)) = ( (2(2661)),
zdzo 1(5) = z(w))) (n)
where (() = 0 and (d) zdzi1e) for all & E X.
For general 0, have
z (E)
z (e) dz (21118)),
z z(W))Fantte n70wi X EE AM 2ell z
where E. C'(ja(Ea); vs) for 8 < l, and Z(Wh) (M)((E) a (11) random variable distributed
azdh((e)
a(Wt)Fan(is)) (W3) ((w))) +
Σ
z" E
O2W3.
vEV: WODEY
For all t 2 0, the limit of the gradient norm is given by
1/2
|| xi ( (i2 ( &) +
| B[Z(dn)) (dh))
11812111111114 11 0 11 n
(af
Taf
(E.3)
where Xt- ('(f(())), vr). Pinn == n1122d
2d: 0z 02
n2-22.03 for 1 E [] and 02 |1
nl 12222,+1, and where 82 ffu ||
if and only if d 1/2 and 82 Biacta if and only if
dr+1 1/2, while /2 [.
if and only if 2d, = 1 + T(I 1) 2
and 02 |o
and only if 2d, > +++1 I(I > 1) 2 22im(b)+1.
For the last-layer weight perturbations (for 0g 2 Ör+1, else z Wh,t1 zow7t1 have
z zud W/!+ + 20 So
zöw!
P 8t zo!
|8|1
Note that År cancels itself out and we purely get perturbation in distribution b² scaled to have standard deviation p.
For all 2 0 and E [1, L], we have
zM- + 0,
ztt= = - + 8
where for 1 -= 1,
ggnt(e) = (Ro(ET E) zan!
||
29

---

If 0 = 0, then
2 - )((2))2zn)
otherwise 0, 1 and
z -(2ht) (2(21))
For / 2 2, we have
11-1)/1 Z 2
Σ Re 1 " z00z -1 1z AMM
==1
Xt
E " z 1zah/
+ +ppwi
/ ||
where
11-11/1
(w 111
and w
Oyt 4,
and Z
has the decomposi-
6,
d:
tion
Z" E aZ
2 2wbbe
-2
vEV: (W.) NE
The perturbed output function has the limit ft = ft + 8fc with
ZZ 2 :. 1l 5w zt$
so that we can define E. C'( Mt) equivalently de L' C'(0L+10 W 2et1,80). For the SAM gradients we have
7 7 Ad
Z W + +0(1+1)/ Zö!!
7, AdbisAMa =
7da'sAM. &6 (2h1))
Zah'sAM z(W) anlaare m70wi Σ Xe- EE E|Z (2(ZaNi ull AaaMaa--dhhhsaM. 2
Elz ah! aatbs 1z Xt
+ ppw:
|1:1|
where (WW))aalsam.. is given by
0zah
Z(W)) abba)aala....... = ad
dilsame +
Σ
Z E
NEV: WODEV
Now SAM's (pre-)activation updates given by zz! zbb (Z0 + + 20hh),
zo1 z + (Zx() + 200z)
with, for / E |2, L). (E)
(ET - - 155-1 1) Z aMC
z
ا/-11/1 (2Wodut np Σ X 2 12
n / /xt - a z ' zddisaam
where 01-1)/11 == 1/11/ Owe/1 = Ow:00 and Z W 1 has the decomposition
2 E 8z62!
Z -2W5sx!-1
Σ
02
vEV: (W.) TEV
30

---

If 6: = 0, then
Zdo ('(zhi- 1 )
otherwise 0, and
7dd (z - 2 (2
The last-layer SAM weight update is given by
z = 2 11
+ /
+ zow
with Z5W --i - 1
Fort > O, the SAM function update is given by
f. - fo + 8ʃ1 + ..... 85
with 8fc =0 B[zów 2 0/ BIz 200]
E.3
Concluding the proof of all main results
After writing out the NETOR program and its limit, as well as tracking all scalings, the main results stated in Appendix D all follow from the Tensor Program Master Theorem and from the characterization results in Yang and Hu (2021) the following way.
Formally Yang and Hu (2021) show feature learning for SGD with small enough learning rate rj > by proving E (5a) f at 0, and they show that learning does not occur in the kernel regime by showing 02, + 0, hence fi fo is not linear in 17.
Both (Z(Z and fi are defined via NESOR T commputations and can be written as a composition of additions, multiplications, the expectation operator, applications of and ø', overall applications of infinitely differentiable,
functions to (Gaussian) random variables, and Consequently 3(Z=f(6o) and fi are infinitely often differentiable where differentiating the expectation operator is covered in Yang and Hu (2021, Lemma H.39). Since Yang and Hu (2021) ) cover the p 0, their proofs immediately show the correctness of the derived scalings for SAM as long as > and > are chosen small enough. Both the gradient evaluation for the perturbation
function of both 17 and
well as the gradient evaluation for the updates stay arbitrarily close to those of SGD if p > 0 is chosen small enough. The conditions for stability, nontriviality, feature learning, perturbation nontriviality and effective perturbations follow from considering the respective scaling.
E.3.1 Proof of Theorem D.2
A bed- padaameatttriiiionnn is stable if and only if all scalings in the Tensor Program have the limit i E 1), 1}, where 0 is required for activations at initialization (for which nothing changes compared to SGD). Potential cancellations are taken care of for sufficiently small r > and p > by the argument above. Now collecting all constraints that are already stated in the Tensor Program formulation at the respective step concludes the proof.
E.3.2 Proof of Theorem D.3
stable bed-parraettriizatinn is nontrivial if and only if f. = O(1) and only if 0/ 1 or
E.3.3
Proof Theorem D.4
stable aadmmerramttinntioon is feature learning in layer if and only if the feature update scaling
= 1 where
T
min CL+1. d++1 + 1/2) + min (cm I1( # 1)).
Hence stable bed- parametrriianno is feature learning in layer if and only if 0.
31

---

Since for all 1 S l2, it holds that rı. 27,, 2 O. we get the equivalence for any to [L]: A stable bed- parametrizaaio is feature learning in layer io and only if feature learning layer for all I 2 lo if and only if r%-- = 0.
E.3.4 Proof of Theorem D.6
Given a stable bodd-parametrization perturbation triviality is fulfilled if and only if 0% f++1 0 and
0. where +/ 12/2-d211 and , n0v0 = n - 1:0211)- hence if and only if d > 1/2 and min cI+1) +
In that case. fi = ft. but ft may still be affected by non mon-aiiihing SAM perturbations in 8W!++1 and oz Only when all SAM perturbations vanish are we effectively only using SGD. By definition, the perturbation scale in the I-th layer vanishes if and only if 0 O, where 0 n - with min(br+1, CL+1) + 1/2 + min
1(m * 1)). hence if and only if Fz > 0. Since for all [4] if and only if > Similarly, for any reference layer lo [Z], we get 8j for all I s to if and only if rio 0. In words, for any la E [L], we have vanishing perturbations in layer to if and only if we have vanishing perturbations until layer if and only if Ft. > 0.
Fz = for all / < L, we get 0:
Altogether, a stable bed- parametrrriiino has vanishing perturbations if and only if 7 > d 1/2 and min(bL+11, CL+1) - F > 1. This case reduces to the results in Yang and Hu (2021) in the limit. Since stability requires 1 2 and 2 O, we can rewrite the equivalence conditions as 1/2 and f max(0, br+1).
dr+11
E.3.5 Proof of Theorem D.8
Recall 0wi := n- (d+td) 0y dyr := (d-((d+)))) By and, for the last layer f := n (d+d2
As opposed to perturbation nontriviality, are not only interested in 0; max Ow:) max'n- 6W... - 1, but in non-vanishing contribution of the perturbations in layer i... Bwr or, for the last layer, 0z+1 1.
E.3.6 Proof of Theorem D.9
The limit of the gradient norm is defined as a NE@OR T program scalar (E.3). Note that for bL+1 1/2, the last-layer scaling strictly dominates all other scalings leading to the simplified gradient norm formula.
Now consider an arbitrary stable choice of layerwise initialization variances {by (bt]re(r++ I] and learning rates {ci)EELL+1] To fulfill the gradient norm constraints (D.1), we have to choose di = C = 1/2 for all / E [L+ 1], because stability requires min(br+11, function perturbations requires d 1/2, where d > 1/2 yields vanishing perturbations and d 1/2 yields effective last-layer SAM through the term 5w.++17f After choosing d 1/2, get 7 2 min(b1+1, CL+11 2 1/2 > 0 which implies vanishing perturbations in all hidden layers.
2 1/2. Now stability of the output
E.3.7 Proof of Proposition D.10
To achieve non-vanishing gradient norm contribution of the last layer in (D.1), we need to choose d 1/2, which requires d 2 1/2 for stability of the output function perturbations. Achieving non-vanishing gradient norm contributions of all layers requires d, 1/2 min(br+1: CL+1) and di = 1- min(br. (b CL+1) for / e [2, LJ. which results in F = d 2 1/2 > 0 which implies vanishing perturbations all hidden layers.
E.3.8 Proof of Theorem D.11
Given a stable bedd-parametrization,, we know d+ d 2 1, so that the feature learning constraint r is not affected by any stable choice of U [di The maximal stable choice of layerwise initialization variances {bi (e(r++)) I] and learning rates (aihett-4) that constitute pP is therefore unaffected by the perturbation scalings d U {di (di)e(LL+))
32

---

Stability of the output function perturbations requires PL+1 + + 2 1. Hence if bL+1 < 1, then i F21-ba O, which implies vanishing perturbations in all hidden layers.
In cv = 1, but effective perturbations in all layers can be achieved more generally for ev 1. Choosing dj 1/2 ev saturates the gradient norm constraint (D.1). To reach effective perturbations already in the first layer 11 cv + d d1 0, need d -1/2. For perturbation stability and last-layer effective perturbations, we need d + d = which requires d = 3/2. Achieving perturbation stability and effective perturbations in all hidden layers requires Ow = which is equivalent to L). we therefore need di 3/2 cp. This choice of
From now on consider bL+1 2 1. Recall ev == min(bL+1,
cv + d di - I(L * 1) 0. For (dihe achieves effective perturbations in all layers.
To show uniqueness we iterate through all possibilities of saturating the norm bound constraint (D.1). We have considered the cases dL+1 1/2 in (b) leading to vanishing perturbations in all hidden layers and di 1/2 ev in (c) with only one choice for effective perturbations in all layers. Lastly consider di 1 - cv for
[2, L] for non-vanishing gradient contribution of the hidden layers. Note that all hidden layers play the same role in all relevant constraints. Effective perturbations in any hidden layer / [2, L] requires 0wr for which need d O. But then, as dj 2 1/2 ov. it holds that F1 2 1/2 implying vanishing perturbations in the first layer. This shows the uniqueness of (1).
For the gradient norm statements, note that the gradient norm ||ell can be written as NE@OR T computation rule (E.2) where the layer scalings in this parameterization are (1) for the input layer, (n1-122 for hidden layers and ((n=1) for the output layer. Now the Tensor Program master theorem immediately implies the result.
E.3.9 Proof of Proposition D.12
Perturbation nontriviality with respect to any hidden layer equivalent to = 0. Since min(br+1: CL+1) < we get min(br+1 (1+1)++7 + < 1. Since stability requires min(br+1, CL+1) F2 1, we get (ii(bb+11,(41)) + == 1, which implies perturbation nontriviality with respect to the output.
E.3.10 Proof of Proposition D.13
The constraint is the same constraint as in Theorem D.8, which implies effective perturbations in the first layer. Now F I fi = -0 implies perturbation nontriviality in all hidden layers due to Theorem D.6.
E.4 Analytic expression of the features after first SAM update
Below we state the analytic expression of the first SAM update, but leave closer analysis of its fine-grained dynamies in comparison SGD to future work. Before looking into the effective perturbation regime, we restate Lemma H.37 in Yang and Hu (2021) with more detailed proof.
First, we define é [4] as the unique index that satisfies Or
Og 0 > >> 01. In
words, E is the first layer in which feature learning occurs. Analogously, define & [L] as the unique index that satisfies
0r 0,
0z
2
81
0,
0r.
Lemma E.2 (Features after first SGD step). Defining Z = zhd 7((1) (((Z)))()))))))))) 1, c vi1(1) (o((2) )')')))) we have
7dt E and
7
76-1 71 Z
and, for all 1 2 e,
z 2/ + 115 11se2 Wosat n18 +ngtzadoo ('
where 3 is defined recursively by
= (g(((1) (1) = -X0- + 3² 3
with g(-1= = 0. Note that (0(0) < 0 for all l 2 l.
33

---

By the defining wndinite-width equations, assuming W = 1 (so minimal stable choice of a).
Proof.
z Zg + ((e-1)) W -xxnn" zdd ø' (Z)).
Atl f, we get 8 (e- 1)/4 0, whereas for I l we get 11)
1, which results in (e-1)/ lise.
Now, for l > l, the second term decomposes into Z and
2 Zuhh 87
OZ(
Since by induction hypothesis,
zdu z ( - 26 = ( + 11se2Wosz + nntl n 21(z)-1) -2)
where Zdrh (((W))" an's is the only dependence 2(W3) ano, get
a (2)))) and o' 3n(z
d' o'(zz-1).
Plugging the derivative back into the defining equation and noticing that Zdub z o' (Z)) concludes the proof.
An analogous analysis for the perturbation at initialization shows.
Lemma E.3 (Feature perturbation initialization). The perturbation trivial layers fulfill
2h6 = 200
hdd = zzb,
and, for all 1 2 E
zho = znb + 11 12222We W
+ p zdda o' (Z²b).
where 3r independent of is defined recursively by
3' = =p)(p) = Xo E[6(Z* )o(Z*b")) "1+ + 3 ''o'' )0' (2hbz))))
|VToll
with 38-1 O. Note that (1(0) > 0 for all 1 2
Remark E.4. If ē = 1, in the definition of 3' replace E(0(Z%)'))(2 " by ST...
Now are ready to state the closed form expression for the first SAM update.
Lemma E.S (Features after first SAM update). Defining Z = 2hd and 23 == 2 we have
- 20-1 2 Z
and, all L 1 t
z Z/ + Tise2 Wooot + n/ Zdo's Zdosamoo ('())) (Z/) + 179 Zuti.
where B is defined recursively by
3¹ == = =
(u (z)-1))
with gt 0, and 7¹ 7¹ (n) recursively defined by
7/ := 30 (4
)0' )0((21-1)
with 78-1 gt = 0.
Remark E.6. If l = 1, in the definition of /3) replace Elo( 1(o(21-106(2-111 by (5)
34

---

Proof. By the defining infinitewwidth equations, for / 2 E, assuming /w1/1 (so minimal stable choice of ct).
Z = Z/ + 0 Wooz MioBlo(2) -0(22)0))0111111 Z SAM.O (' (Z)).
(E.4)
At / E, we get 6 - 1)/e = and /we/e= = 1, whereas for we get 8 (1-1)/t 11/1 = and 9 W / = (under minimal stable choice of et), which results in (1--)//1 - -|s--- Now, for 1 > the second term decomposes into Z and W For the rest of the proof it remains to analyse 2
Since by induction hypothesis,
zoo
)( --(10-1
+11 @ (z) + -11 2. 2Wobof + +nol-1z our
saMa a' (Z0-1) 1p--111111 zanb ( 6(Z0-11)
where Zanaaaso = z(W)) (al))ad + pēw: to E Zdt zdl'' ydlbssas..) 1
with the second term
|V
independent of (W and by Lemma E.3 we know 2-1
Z011
111 180
p Ø (Z)-1), where only the last term with Zdzo z(W))" dh' influences 2 we get
Z aZóm
82
2 W
SA
(E.5)
W2 dh!
82 (w:) FFdd))MM
with
az -1
= =6' =Z 1m3 31-1 0' (Z--1) )
32 dh'gAM.
and, using
2'(Z) - (( dab dbl (Z)
(Z--1) 1) yields
0Z
1761 2u 1-1 0" (Z)-1111 026-1
0'(2-11)
+ 11 (2)-1)
(2(W))
02 02:Wy:Taby
s' 1 (n/gt-1 1i (2 1 (z m 06 62(2(-1) 1))
Plugging Eq. (E.5) back into the defining equation (E.4) and noticing that 7dhlsAMa Ad o' (()) as well as d Zddi o' (Z)) concludes the proof.
F
Generalizations and further perturbation scaling considerations
F1
Overview over choices of d, and d
Since for some combinations of architectures and datasets it turns out that performing SAM on subset of layers performs better than effective perturbations in all layers (Müller et al., 2024), we would like to know how to choose d and di to adjust which layers should be effectively perturbed and which should have vanishing weight perturbations. In practice, simply set all perturbations that should vanish to 0 by design, and use the global scaling d and relative scalings di from jaP² for the perturbed layers. This section instead interested in a complete characterization of all possible choices of (di )e(L++ and d. The heuristic derivation only requires the gradient norm constraints (D.1) and the perturbation stability constraints that require Sw1 = O(1) and åw = O(n for I >>1 given by
c d
if l is input-like, if I is hidden-like, / output-like,
dj cv
(F.1)
1-d d
where a layer effectively perturbed if and only if equality holds in the respective perturbation stability inequality. This heuristic claim yields the characterization of all phases of the choices of
35

---

Effective perturbations possible
Gradient norm may be dominated by hidden-like
input-like hidden- -ike
output- like
input-like
output-like
d= 112
X
d E (-1/2.0)
X
X
X
d= 0
X
XXX
de 1/2)
X
X
d= 1/2
X
d> 1/2
X
X
Table F.T: (Characterization of perturbation scalings) Overview over the regimes of all possible choices of d and dj. A layer is effectively perturbed if and only di satisfies (F.1). At least one layer must satisfy equality in its gradient norm constraint (D.1). This table summarizes which layers can exhibit effective perturbations, and which may dominate the gradient norm, given a choice of d. The choice d < -1/2 results in perturbation blowup F < 0. At the critical d -1/2 1/2 (respectively, d = 0; d = 1/2) a input-like (respectively hidden-like, output-like) layer is effectively perturbed if and only it dominates the gradient norm. Consequently d 1122 implies effective perturbations in at least one input-like layer.
perturbation scalings d and di in Table F.1 and allows us to formulate a simple rule of how to choose d and di given the information which layers should be effectively perturbed, and which should have vanishing weight perturbations.
Choice of perturbation scaling from list of layers to effectively perturb. We denote the set of all layers by L. whereas the subset of layers, which we want to effectively perturb, is denoted by CSAM S C.
1. If there exists an input-like layer 1 E LSAM. set d -1/2. Input-like layers are effectively perturbed and only if d, 1/2 <v. Hidden-like (respectively, output-like) layers are effectively perturbed if and only if d; 3/2-cx (respectively, di = 3/2). For all layers that have vanishing weight perturbations, do not perturb these weights or choose di > 1/2 - cv for input-like, di 3/2 cv for hidden-like and di > 3/2 for output-like layers. 2. If all input-like layers should have vanishing weight perturbations but there exists a hidden- like layer E LSAM. set d O. Hidden-like layers are effectively perturbed if and only if di 1 cv. Output-like layers are effectively perturbed and only d = 1. For all layers that have vanishing weight perturbations, do not perturb these weights, set di cv for input-like, d > 1 - cv for hidden-like and dj > 1 for output-like layers (as required by the perturbation stability and gradient norm constraints).
3. If both all input-like and all hidden-like layers have vanishing weight perturbations, but there exists some output-like layer / E LSAM. then set d 1/2. Output-like layers are effectively perturbed if and only if d; = 1/2. For all layers that have vanishing weight perturbations, do not perturb these weights or set di 2 1/2 - ev for input-like, di 2 1- cv for hidden-like and d; > 2 for output-like layers (as required by the perturbation stability and gradient norm constraints).
4. If CSAM = 0, then set d > 1/2 or simply perform SGD.
Example E.1 (First-layer-only effective perturbations). Instead of simply using the rule set above, we derive the necessary choice of perturbation scaling from the scaling equalities and the norm constraints (D. To achieve first-layer effective perturbations, but trivial weight perturbations in all
other layers, we need 6y = 1 and u 0, for which we will choose 6 = -1, This requires setting
d1 -cv + d),
di 2 cv d,
dL+1 2 - d,
where one of the constraints (D.1) has to be fulfilled. Plugging the above dy-choices into (D.1) results in the constraints d < 1122, d < 1, d < 3/2, hence choose d = 1/22 so that only the first layer contributes non-vanishingly to the gradient norm. Note that 7 and output perturbation nontriviality holds if and only if min(br,-1, c1.+1) in Appendix H.2 to show that propagating perturbations from early layers are not enough to inherit SAM's ss inductive bias that leads to improved generalization performance.
(as in jP). We apply this perturbation scaling
36

---

F.2
Other ways to introduce layerwise perturbation scaling
Before presenting alternative ways how layerwise perturbation scaling could be accomplished, let us collect desirable properties that a definition should fulfill:
Layerwise perturbation scaling should enable stable, effective perturbations in every layer. The perturbation step should require at most one additional forward and backward pass in each update step.
The adapted optimization algorithm should recover the original (SAM) algorithm when not using layerwise perturbation scaling.
We start with the simplest case where the perturbations are normalized in each layer separately.
Remark F.2 (SAM with layerwise gradient normalization). For (SAM) with layerwise gradient normalization of the perturbations
=' WwiL///V LII,
(LN)
where || 1. may denote the Frobenius or the spectral norm (equivalent under limited perturbation batch size), the spectral scaling condition (+) immediately yields the correct layerwise perturbation
scaling 1 e fftt_oouutt ou..ininif.n........
However, in practice, perturbations are usually globally normalized across layers, according to the GitHub repositories provided by Foret et al. (2021); Samuel (2022); Kwon et al. (2021); An- driushchenko and Flammarion (2022); Müller et al. (2024). Preliminary ablations in Appendix H.5 suggest that layer-coupled SAM with global normalization slightly outperforms SAM with layerwise gradient normalization. As our goal this paper is study (SAM) as applied in practice, we consider SAM with joint gradient normalization.
A first alternative to Definition 4 could scale perturbations after the joint gradient normalization.
Opposed Definition 4, for this variant the perturbation norm, i.e. the radius of the adversarial ascent ball, is not guaranteed be pn -d but p
p(EiE(L+1) p?)1/2 The correct perturbation scaling
for this version more immediately follows from the condition that perturbations scale like updates. Remark F.3 (Layerwise perturbation scaling after joint gradient normalization). For (SAM) with joint gradient normalization of the perturbations
VwiC
e' = Pt Vw
the correct perturbation scaling in jP is given by Pt - (nn2/2 ffa_oou/faal
To understand this scaling rule, note that for bL+1 1/2 (such in μP)), the last layer always dominates the gradient norm (see Eq. (E.2) for the TP argument), resulting in the scaling
V = C'(f(E)), 11)l==1 = (22
Thus, compared to SAM without gradient normalization (Appendix F3), |/ W always con- tributes the scaling n 1/2 Noting that perturbations should scale like updates, and updates receive the layerwise learning rates m O(fan_out/ coutt/fan_ir //f//annnnnn in) concludes the derivation.
In Definition 4, accept the additional layer-coupling complications that the layerwise gradient scaling before the joint gradient normalization entails in order to analytically control the perturbation radius to pm -d To simplify the analysis as much as possible, we will first ensure width-indeeendeenee of the normalization, so that the layerwise perturbation scaling is not affected by the normalization term. Then, layerwise perturbations should be scaled like updates.
Another alternative layerwise perturbation scaling as in Definition is motivated by the observation, that in paP2 with Definition 4, only the first layer dominates the joint gradient norm (Theorem D.11) To let all layers contribute indehiindeeendently to the joint gradient norm, we can introduce even more hyperparameters (with limited benefit) by decoupling the numerator and denominator scalings the perturbation. Opposed to Definition 4, the perturbation norm is again not analytically set with the choice of p. but may be smaller. Empirically, we do not observe performance differences due to denominator contribution scaling (Appendix H.4). This is the perturbation scaling we implement for ViTs (see Algorithm for details).
37

---

Remark F.4 (SAM with decoupled perturbation numerator and denominator scaling). For (SAM) with perturbations
V W ||||
L++1
e' pn
with ||al|
Σ
n 2d
V W
(DP)
with layerwise perturbation radii p n and separate gradient norm scaling n -d.
In all alternatives, nontrivial layerwise perturbation scaling is necessary for effective perturbations in every layer, which necessarily changes the direction away from the original gradient direction. Such layerwise gradient rescaling can also be achieved by adapting the architecture with width- dependent weight multipliers. The multipliers (a-pP22 achieve effective perturbations without layerwise perturbation scaling such that all layers contribute non con-vnnishingly to the joint gradient norm. They rescale the gradients equivalently to (DP) when scaling all denominator terms to be width-independent See Appendix F.6 for all details about weight multipliers.
Adapting the TP-hased analysis. Our based analysis covers all of the above perturbation scaling alternatives with minor adjustments. We just have to replace the normalized TP scalar (E.2). If we want to express [| we just drop all perturbation scaling terms d: For the examples of (LN) and (DP), we replace (E.2) in each layer separately by the normalized TP scalars,
Vw Coll = Xx ((d))))
(dhh)T (idT)dl dh! (ET(E))
1/2
with scaling 811,,11 = n 1/200 for the first layer, where By is overloaded denote 0y n br.+1 in the first step and 0r = n min(b++1,,4441)) in all later steps (in puP, fv = n always).
(dhh)) dnd (mt -= 1-1
1/2
V Li
Xt
n
7
with scaling 0/ now for all hidden layers E [2, L]. and
(a (at) 1/2
IV Cell = Xt
with scaling vn for the output layer, with respective normalized limits
Xe (()))) ((T[ 2/) (E (2)) "E(Z "
(
where Xr (f. yt).
The adapted scalings can then be tracked as before derive the maximal stable layerwise perturbation scaling. Consider for example input layers in (LN). In jP, we know w/ Lell O (1/2) and V Wwi Le O entrywise. Effective perturbations are achieved with e¹ = O(1), so choose as expected from (*). Proceed similarly for all layers and perturbation scaling variants.
n =
F33 Extension to SAM without gradient normalization
Andriushchenko and Flammarion (2022) and Andriushchenko et al. (2023a) consider the SAM update without normalizing the gradient in the adversarial ascent. The corresponding update rule is given by
W. W. - v, (E, We + prr, (e), ve);
1t = V. W.).
The NE@ORT program for this update rule with arbitrary of adapted from the above derivation. Just note that the gradient norm appears in an equation if and only if the perturbation radius pn pn-4 appears. Without dividing by ||ve]]. the parameter d becomes superfluous. Simply set d = 0 and remove the gradient norm constraints (D.1) to arrive at the NEmORT program and bed-constraints for the update rule without gradient normalization.
Wt) is also easily
Perturbation scaling di plays a similar role as learning rate scaling e as both scale similar gradients. We get effective perturbations in the 1-th layer from the equation di + min(br-1,, CL+1) q
38

---

min(br+1: CL+1, 1 in pP, which yields di qi for all I E [L] (since 1=+p 1 for stability). In particular, in jP,, the correct layerwise perturbation scaling of unnormalized gradients is given by the rule fan out or the squared weight (update) spectral norm || (Yang et al., 2023a), which could be efficiently approximately tracked with running power iteration.
fan in
Note that Dai et al. (2024) argue that the normalizing the gradients for the perturbation is crucial (in standard parametrization) due to stabilizing effect and an enhanced drift along manifolds of minima. Monzio Compagnoni et al. (2023) find that unnormalized SAM gets stuck around saddles while SAM Hessian-inducedd slowly escapes through additional Hessian- induced noise. This suggests that the additional effort of analysing the original SAM update rule with gradient normalization is necessary for practically useful theory. From this paper's point of view, the gradient normalization may be adding stability via the n 1/2 contribution which allows scale down p less aggressively in practice.
E4 Extension to Adaptive SAM
Adaptive SAM (ASAM) (Kwon et al., 2021) is motivated by sharpness definition that is invariant to parameter rescaling operators that leave the output function invariant, and can provide a further improvement over SAM of 0.5% to 1%, depending on the considered vision dataset and model (Müller et al., 2024). Here we consider the two examples of elementwise rescaling operators (with
p = 2) and layerwise rescaling operators (with p= 2), which are the best performing SAM variant in most settings in Müller et al. (2024).
Proposition F.S. Neather elementwise ASAM, which performs (SAM) but using the perturbation rule (F.4), nor layerwise ASAM, which performs (SAM) but using the perfurbation rule (F.6), can be written as a NESOR program.
Proof sketch. Elementwise ASAM requires an elementwise multiplication of matrices, and layerwise ASAM requires calculating the Frobenius norm of a matrix. A NonLin operation only takes vectors as arguments, so NEgOR calculations with a matrix require its multiplication with vector. But then a single coordinate of the resulting vector contains a mixture of an entire row of that matrix. Since we are only allowed to define random vectors and matrices, and the NEgOR T master theorem states that coordinates of NE@OR T vectors behave iid-like, this mixture cannot be disentangled by choosing a structured vector. Hence, already initialization, the square of individual entries/the Frobenius norm of random matrix cannot be exactly recovered by function of matrix-vector products with NE@OR vectors.
Although ASAM is not formally covered by theory, we still expect that the ASAM perturbations are correlated with the gradient and therefore with the incoming activations, so that heuristically still expect LLN-like behaviour and apply our scaling condition. If the perturbation rules still behave LLN-like, then Table summarizes which layers are effectively perturbed under global scaling and provides the unique maximal perturbation scalings for all considered SAM variants. The correct perturbation scaling in μP for other perturbation rules that behave LLN-like can always be derived following the same steps:
1. In μPP, always holds that
(1)
1= 1,
(0(0v) = O(n i L,
w/_ =
((n-1/2)
lE [2, L.J,
and V (0(1= (0(1)
[= L 1.
0((-!) 1= L 11,
(F.2)
2. Assuming the normalization term in the denominator is scaled to O(1), track the layerwise scalings of the numerator. Maximal stable perturbations are always achieved with
(0(1) 5w! (( 1>
1=1,
(F3)
This yields constraints for achieving maximal stable perturbations in each layer.
3. Now replace the norm constraints ().1) by tracking the scalings of each layer's contribution to the update rule's 's total normalization term.
4. To ensure normalization term scaling O(L), iterate through the layers 4:
39

---

(a) choose di to satisfy its norm constraint,
(b) choose d to induce maximal stable perturbations in that layer,
(c) choose all other de, l' * 1, minimal to both satisfy its norm constraint well as
(0
1=1,
perturbation stability 8w? 0 1>1.
5. From the above confi suniigurations, choose the unique one that yields maximal stable perturba- tions in all layers.
F.41 Elementwise ASAM
we want to be invariant to elementwise rescaling operators TV, (x) = | z where a, W Rm and @ denotes elementwise multiplication, the resulting ASAM perturbation rule (where introduce (layer-wise) perturbation scalings {d} U [(d]u(iihhr {dihee(t+)) replaces (LP) and is given by
āw = pm d [W' @ |W'| O W(iE( W))
(F.4)
Velem
ASAM
with normalization
L+1
V
elem ASAMI =
Σ n
W (.)
where the absolute values [W'| are computed and multiplied elementwise. To find the correct perturbation scalings, we track the typical elementwise scaling of each quantity as before.
Elementwise ASAM in J.P. In jP,, the layerwise weights and gradients scale as (F.2). For ASAMl = O(), we therefore replace the constraints (D.1) by the constraints
di 2 1/2 cv, for l E [4],
dr.+1 2 11/2,
(F.5)
where we can choose {dhEEL+11 to achieve equality at least one constraint to achieve || ASAMrl @(1)
O -x) O(n (( ((nin)=)
1=1,
The layerwise perturbations scale as ów: = n -d
I l [2, L].
1= +++1
Stable and nontrivial perturbations each layer are achieved under condition (F.3), which induces the constraints for optimal layerwise perturbation scaling
d d di = -ev, for E [L],
d d+ d++++1== = -1.
Irrespective which of the above norm constraints (F.5) we satisfy, we need d = 112 to achieve optimal layerwise perturbation scaling. Hence d = d -1/2 and d, 1/2 cp for / E [4]
the unique choice of {d} U {di ]e(t.+11 modulo norm scaling equivalence that achieves O(1)
perturbation scaling in all layers. With this choice all layers contribute non-vanisee non-vanishingly the gradient
norm. In μP cy = 1, so that d, = -1/2 for all [L + 1). so that ASAM does not require layerwise rescaling of the gradients, but upscaling the perturbation by n 1/12 to achieve nontrivial perturbations in any layer. This may explain why ASAM often outperforms SAM in large models: By only requiring global scaling, ASAM achieves maximal stable perturbations in all layers the perturbation radius is tuned globally at every width.
instead of a global gradient norm A Il. one would want normalize in each layer separately
with V eeemm. ASAM || == n d ||W @ Vwi W.), Irr,,, the layerwise perturbation scalings
become ōw? n-d (n 1/2)
1= 1,
Again, to achieve maximal stable perturbations in all layers need d 1/2 and layerwise adaptation of the gradient norm.
10 3/2) (3/2)
1> 1.
'There must exist such a choice of {d. he[L-1] and d, because (di he(c=1] allow set any relative scalings between layers and d determines the global scaling, which overall allows to set all possible layerwise scalings. Any deviation from the choice that achieves effective perturbations in all layers either results in blowup or vanishing effect of the weight perturbation some layer. This shows uniqueness.
40

---

F.4.2 Layerwise ASAM
ASAM with layerwise rescaling as in Müller et al. (2024) employs the layerwise transformations
Ty (a) = || This ASAM perturbation rule replaces (LP) and is given by
kw! = pm mndd di | 11 W (C W.). ve)
(F.6)
V lager
ASAM
with normalization
L+1
layer ASAM
||
Σ 2m1
((
1 1,
Layerwise ASAM in p.P jaP, we have ||
((2)12)
IE [2, LJ, 1=1+1.
Hence, the norm constraints (D.1) are replaced by di 2 ov. ov, di 2 3/2 cv for l E [2, L],
dLT1 2 0.
(n di nl)) (n din-dint (nov)
1= 1,
The scale of the perturbation numerator now scales as 8w - n -d
lE 22 L],
Onn do--in
I = L+1.
In jal, achieving maximal stable perturbations (F.3) is therefore equivalent satisfying the constraints
d+ O, d+ = for / [], d d 0.
Now we can simultaneously satisfy the first- and last-layer norm constraints with d1 = 0 and d O, while achieving effective perturbations in all layers with and di = 1. Satisfying the norm constraint in the hidden layers with di 1/2 would imply vanishing perturbations in the first and last layer (by requiring d 2 1/2).
F.5 Representing general architectures and adaptive optimizers as Tensor Programs
Here, we lay out explicitly how to write some of the building blocks ResNets and ViTs a Tensor Program and provide further scaling considerations. According to Yang and Hu (2021), it is straightforward to generalize scaling conditions that induce feature learning in MLPs to these other common neural network building blocks. Since perturbations should always scale like updates, the conditions for stable feature learning and those for stable effective perturbations are analogous.
One potential complication in the case of SAM would be a contribution to the joint gradient normal- ization |/ll|| that differs from the classical input, hidden or output layer contribution. But we will see that these contributions do not differ for any of the considered layer types.
Layernorm. The Layernorm operation is defined
7t x.. V/
h
B/-
0
where E > 0 is small positive constant, N/. 3, are learnable parameters and 1/ = Σn (zt); an Avg operation as in Yang and Littwin (2023, Def. 2.6.1) and of / 1 Σi (24 1/22 is composition of Nonlin, Avg and Nonlin. The parameters 7/. St can be seen as input weights to the input 1. They should be initialized % 1 and B 0. In the forward pass, the layernorm preserves stability h,t1 O(1) when o. 3 @(1) except for the Lebesgue nullset of learning rates for which they exactly cancel each other out. Recall the notation dz = 0-10f /0z, where fg n for some C' IR denotes the width-dependent scaling. The derivatives
d7t = dht (+1 1 v
d33 = 1+1
of +
41

---

These gradients coincide both in shape and scaling with the scaling we expect for an input layer, resulting the same gradient spectral//robenius norm scaling, Continuing the backward pass, using
Bo #1 -u
get
Out
not
----) 1)
a - v4 do!
dzl =
dhitlng 78
n (0 dxl
0 +
- (o) 1.-1y v1)
dh; dhl+1 1/
(I 117)
of +
n
(o) + e)² not
which preserves the order long 7o (1), since xa ((1,, we know 1/y, =
Note that Layernorm removes the necessity to avoid blowup in the activations x; in the forward pass (ignoring potential numerical issues), and always rescales to (m,ax(7,, 8))) However, in the backward pass, a scaling z = O(n"), with e > O, results in dx O (n-
+ (1)) hence vanishing
gradients. The gradients would only stabilize if (h()) = @(h)). ). but no popular activation function has scale equívariant derivative. Yang (2019) shows how to write Batchnorm and Average Pooling as Tensor Program.
Convolutional layers can be seen as a collection of dense weight matrices where width corresponds the number of channels (Yang, 2019). With kernel positions ker, input channels [n'] and output channels [n/+1), the weights of a stride-1 convolution are given by Wa (eee so that for each i E ker, W pn'+'xx xn' is dense matrix. With {at heepeeeesssccoo ae]le] the convolution operation is given by
Convolutions.
(W' a)aa
E
W
jo.3 143.8
Biis+iEpos!
which performs MatMul and Avg and where ker,pos pos' are assumed to be of fixed size, For ker of fixed size, convolutional weights scale like hidden layer weight matrices, also in Frobenius norm contributing to ||u:]]-
Residual connections. A residual connection propagates the current activation forward, skipping an arbitrarily complex nonlinear block ft R - " between, where ft can depend on time-dependent parameters like weight matrix. The forward pass can be written as
2 = 2 f(
If ær ((1) for all layers I holds in the model without residual connections, it also holds in the model with residual connections. At fixed depth, f o(1) should be avoided, as it would hold that 2++1 x in the infinite-wine width limit and the layer would be superfluous. The derivative of the activations becomes
dxt-1 = dr. + drl as! a=, -
where the second term stays the same without the residual connection. For the example of f being fully connected layer we get dx!-1 dr² + (W!)² (dx) d'(W(x(-1)) In this example, the derivative with respect to the weights becomes
0f ds! dw! dw!
dxt dey
0fi
(de, @ o' WWaa--))aat- ow!
where the residual connection does not alter the functional dependence on dæ, and ar compared to MLP, but implicitly influences the weight gradient since du; and of are altered. As for the forward pass, the gradient scaling dz; gets stabilized in the backward pass that af' is now allowed au to be vanishing with width. Again, we are not aware of an architecture in which that would be desirable. Since residual connection does not introduce learnable parameters, interferes in [|v, only implicitly through the stabilized gradients in earlier layers, which can contribute -aniisishy to [o even if later layers are wrongly scaled and their scaling is not adapted.
Adam as base optimizer. When using Adam or similar adaptive optimizers as a base optimizer, the learning rate should scale as @(1) for input-like layers and biases, and O(n -1) for hidden and output
42

---

layers (Yang et al., 2022). Yang et al. (2023b) provide proofs for arbitrary optimizers that perform generalized, nonlinear outer products. the example Adam, the update rule can be written as
1/2
= "'mut u
Σ we(uua
where "i; w; are the weights that stem from the moving averages. By using learning rate of n-1 and using the fact that both u and have approximately iid coordinates of order ((1), the law of large numbers yields (1)) updates the form
1 Σ ou ....., u U...................... .
Dg = Ep(u, ................................ u Z" ..... )z
Any other learning rate scaling would either result in blowup or vanishing updates.
Adaptive optimizers have not been used for the ascent/petturaation step. In the descent/update step, nothing changes compared to unperturbed optimization as long ensure stable perturbations.
F6 Influence of width- width-dependent weight multipliers on rode-aarmetteizationns
Our definition of bed- parameterizations is convenient because it purely adapts the learning algorithm but not the architecture. We can also adapt the architecture by using layerwise width-dependent weight multipliers to effectively perturb all layers without any perturbation scaling. The reason that layerwise weight multipliers scale the layerwise gradients. Here, we study how the introduction of weight multipliers affects paddrrramtteriaaionn.
In this section, we consider L-hidden layer MLPs with weight multipliers farhe(t width n N, inputs & Rd,, and with outputs f(E) = n a113 w (E((E where the activations x²({) are defined via the iteration
(E) = n W's
x'(E) = (E)'
A (E) n -
We define abcd- parameterizations in the same way as podepereetttii,iinnnn,, but instead of MLPs we use MLPs with weight multipliers far he(r.+1.
Definition F66 (abad--parametiizatinn)) An poeemeraramttiizatioon {at)E(L++) U (brhee4+1) U (o) U (44 (d]hiett1)) U {d) defines the training of MLP with weight multipliers (arhe((++))
with SAM in the following way:
(a) Initialize weights iid as W ~ (0,
(b) Train the weights using the SAM update rule with layerwise learning rates,
W ++
W rnn V V (Et; W2 + Et) Bt),
with the scaled perturbation e; via layerwise perturbation radii,
E, =
Cr
with
n = (u of ==
Vwi Wt), W.)...) (LP)
W.l.o.g. we set ||aal O(1), which prevents nontrivial width-dependence from the denominator. This imposes the constraints:
d a 2 1/2 cv.
di + a, 2 - d + 2 1/2,
with at least one equality required to hold, where 1 E [2, L]. and where Varf n n- WL+1 O(n- with cv min(b141 + #L++,,C-11 + 2a,+1). The normalization 2 ///] removes one degree of freedom from {di he(r++) via the equivalence (e(hettt+ = (dihe(t++1) iff there exists C' E R such that d' = di + C for all [L+1].
F.6.1 abed-equivalencee classes
Update scalings behave as in SGD. The weight multiplier n -at scales the gradient Vwr ʃ by ai In the following forward pass, another multiplication of the weight updates with n -ai leads the activation update scaling n 201, This can be counteracted by adapting the learning rate scaling, For
43

---

abo- parpaatteriiztionns and SGD training. this induces the layerwise equivalence between parame- terizations with (at, bt,cc) et) or with (a, + 8,, b 8, ct 24.). The extension of all of our results to Adam base optimizer is straightforward, since learning rate scalings and perturbation scalings are decoupled. For Adam, et should be adapted to q --- 0/.
Again, perturbations with joint gradient normalization complicate matters compared to SGD and Adam. Keeping the gradient norm scalings invariant under d; + 0t + 8 would require d, r dN 8, but keeping the activation perturbation scaling invariant would require di ++ dp 20 as for updates. Consequently, an exact equivalence between abed- parameterizations at finite width requires A to be the same for all layers and the conflicting gradient norm in the denominator and perturbation scaling in the numerator be accounted for by adapting the global perturbation scaling d d (together with dj + di 0). In other words, (SAM) with layer-coupling gradient normalization (LP) does not have layerwise analytical equivalence classes at finite width. Below, we provide alternative perturbation rules that resolve these complications and recover layerwise equivalence classes. The following lemma formally states the layer-coupled equivalence relation for the perturbation rule (LP). All proof's are provided at the end of this section.
Lemma F.7 (abed-equivalence classes). Let fr(E) denote the output of a MLP in a stable abed- parameterization with weight multipliers (ei)et(t+) after t steps of training with the SAM update rule with layerwise perturbation scaling (LP) using a fixed sequence of batches and evaluated on input Ş. Then for any 0 E IR and any C E R. f((E) stays fixed for all t and if. for all 1 E [L + 1),
(ar, bi, Ct2 di, d) is reparameterized to (at + 0, bi - 0, 20, di - 0 + C, 0).
Remark F88 Unfinite-width equivalences). In the infinite-widie limit, abcd-patee poodlpprramteeizattonn remain equivalent under (at + 01, by - 0, a 20,, d: 20,, d) layerwise as long as the set of layers that contribute to the gradient norm non-van vaninsiiiinll remains invariant. The gradient norm constraints for ||06|| = O() become
di + 01 2 1/2 12--vv cv,
di a 2 11-10 dr- + 0141 2 1/2,
where F = n # - = -(n with cv min(br+1 + CL+1 + 2al+1) remains invariant under equivalence transformations.
Remark F.9 (SAM with layerwise gradient normalization), As the layer coupling is induced by the joint gradient normalization in the perturbations, layerwise gradient normalization simplifies the anal- ysis. For (S (SAM) with layerwise gradient normalization (LN) of the perturbations global perturbation scaling d is superfluous, and there exist layerwise equivalence classes: For any (8) C R,
(at, by, Qt, di) is equivalent to (a; + 0,, 0g, q 20, di - 6.).
To understand this equivalence, observe that any layerwise gradient scaling cancelled out by the normalization V V LlI. Only the n n- factor from subsequent forward passes has to be counteracted.
Remark F.10 (SAM with decoupled perturbation numerator and denominator scaling). A perturbation rule joint gradient normalization and layerwise equivalence classes can be achieved by introducing even more hyperparameters and decoupling numerator and denominator scalings of each layer. For (SAM) with perturbations (DP) with layerwise perturbation radii p nd= and separate gradient norm scaling n di global perturbation scaling d is superfluous, and there exist layerwise equivalence classes: For any {0) [0])ie[[+1] C IR,
(at, bi, at di, d:) is equivalent to (at 01, , a 20,, d: 26,, di 0:)-
This perturbation rule also allows us to recover an analytical equivalence between trivial weight multipliers a; 0 for all I, and any other weight multipliers.
F.6.2
p2 under non-trivial weight multipliers.
Our goal here to find the weight multipliers that simplify the necessary perturbation scaling for effective perturbations in all layers as much as possible. The non-existence of layerwise equivalence classes in parametrrretssnn from (LP) is not an issue if we are interested in effective perturba- tion properties and recovering ju2 for arbitrary weight multipliers {a])ie[4+1)- as the equivalence breaks due to varying gradient norm contributions, which are inconsequential for achieving effective perturbations.
44

---

As we aim to reproduce pP? we restrict ourselves to the p&P equivalence class of abc- parameterizations. We do not allow layerwise perturbation scaling and are interested in the maximal stable choice of global perturbation scaling pm to at least achieve non-vanishing perturbations in some layers. The following lemma shows even more: The choice
a = -1/2 I(I = 1) + 1/2 -+==++
achieves effective perturbations in all layers with the naive (SAM) update rule with naive perturbation scaling p n², and all layers contribute non-vanishingly to the joint gradient norm. Hence this seems to be a natural choice of weight multipliers for SAM. However, it is in conficct with unit scaling considerations (Blake et al., 2024). Effectively, naive learning rate and perturbation scaling with these multipliers is equivalent to (DP) where all denominator terms are scaled to be width independent, as implemented by Algorithm which resembles our implementation for ViTs. Our ablations in Appendix H.4 suggest that gradient norm contributions have a negligible effect on generalization performance.
Lemma F.11 (Naive perturbation scaling can effectively perturb all layers). Consider an abed- parameterization where {(a by, qi) (e)eE()11) are chosen from the yP equivalence class, and where there is some C R such that di C for all I [I. + 1]. This reduces to training MLP with weight multipliers with (SAM) with global perturbation scaling pm pnd for some d E R. Effective perturbations all layers are achieved and all layers contribute non-vanishingly to the gradient norm if and only if
1 1/2.
a-dd for E L],
G++1 -
Achieving μP2 with the current implementation of the mup- mup-package requires both an adaptation of the architecture and of the learning algorithm, as the following lemma shows. Hence the package is not particularly suited for SAM learning in μP² when the goal is simple perturbation scaling.
Lemma F.12 (Effective perturbations with the mup-package). Consider an abed parameterrizatioonn where {(a bt, ct) he([++1) 1] are chosen from the pP equivalence class, and with the weight multipliers aL+1 = I(I = L + 1) as in the mup-package.
(a) (mup-package global scaling effectively perturbs hidden layers) Under global scaling
d₁ = C, E R. for all [T. +11, maximal stable perturbations are achieved with In this parameterization, hidden layers are effectively perturbed, but input and output layers are not effectively perturbed.
(b) (nP² with the mup-package) Effective perturbations in all layers are achieved with the choice d = di = d = -1/2 and di = 1/2 E [2, L}.
The following lemma covers the general case how to achieve ju2 given arbitrary weight multipliers. Lemma F.13 (nu2 with arbitrary weight multipliers). Consider an abed parameterization where (ar, bi, <t) he)LE(L- are chosen from the μP equivalence class. Then effective perturbations in all layers are achieved with the choice = minge[L+1]] ( 1/2T[[= 1/2I = L+ 1)), and
di = - d- 201,
di - 201, for E L), d 1 d 20
The following lemma shows that weight multipliers that achieve jaP² with naive perturbation scaling under perturbations with layerwise normalization (LN) exactly the same the ones for (LP). Lemma F.14 (LLN) with naive perturbation scaling can effectively perturb all layers). Consider (SAM) with layerwise normalization (LN), Assume {au by, (1, (E) are chosen from the μP equivalence class, and assume there is some C E R such that dj C for all l E [L + 1). Then all layers are effectively perturbed if the multipliers are chosen as
a1 = 112 C, 41 = CC
at+1 = 1/2 1/2-C. -C C.
Proof of Lemma F.14. As derived in Appendix F.7, under a for all [L + 1]. all layers effectively perturbed if and only if d, = 112 I[[ = 1) + 1/2 II[1 L L 1). Now we can exploit the layerwise equivalence relation to enforce d; C in each layer by adapting all az.
Proof of Lemma F.I.. I. In general, the abx-equivalence class of μP, the 1-th layer's gradient norm scaled by n- This induces the generalized gradient norm constraints for V T (1))
d1 + a 1/2, d +a O,
d + a++1 2 1/2.
45

---

Effective perturbations are achieved when pn -d 1odi--n0 V VWLL = *(nn (nn-1))))) which induces the perturbation stability constraints
+ 201 2 1, d+ 20 2 0,
dLt1 + 20 1,
with effective perturbations whenever the equality of the respective layer holds.
Under global scaling, the gradient norm constraints become, for some I,
C + a1 2 -1/2, C+a +
C+ aL+1 2 1/2,
and the conditions for effective perturbations become
d C+C+ 20 2 -1, d+C+ 2a 20
d C 2a 2 1.
As + C' is a common term in all layers, we get the relations dj = 01 + + aL+1 a1 + 1, so that all gradient constraints are simultaneously satisfied with C = -a; and effective perturbations are achieved in all layers with -0t-
Proof of Lemma F.12. Under the choice a; = [[[ L + 1), the gradient norm constraints become
di 2 1/2, di 2 0,
d1t1 2 11/2
and the conditions for effective perturbations become
d d1 2 11
d+
d+ ++d1+1 2 -.
Proof of (a):
Satisfying the gradient norm constraints with global scaling requires d, = 0 for all E [T. + 1). then the minimal stable choice of d is d = 0 which only effectively perturbs hidden layers.
Proof of (b):
The choice 1/2 and d -1/2 saturates the gradient norm constraint and achieves effective
perturbations in the input layer. Then the choice d: = 1/2 and d -1/2 satisfies the gradient norm constraints and achieves effective perturbations in layers.
Proof of Lemma F.7. To understand the influence of weight multipliers on updates and perturbations, first note that under equivalence transformation of all abcd-parameet w.l.o.g from a = 0 for all / E [L + 1). the scalings of h x! and of n Wi remain invariant. This implies that the scalings of Va4f = n #1+1 W V = f @ o' (h") and oftf for all E [L] also remain invariant. Hence the weight gradients, V W f aztirl and V W n-n W (a!-1)T are scaled by n in each layer.
In the following forward pass, we get
ht = n - (w' + "n 1pn w: W:
so that activation/output updates and perturbations of layer are scaled by n -2a.
Again, a complication compared to SGD or Adam arises through the gradient normalization SAM's weight perturbation. If the gradients simply normalized layerwise s p.n the n -trmm from the backward pass cancels out, and only in the forward pass we get a scaling n -ai Hence an exact layerwise equivalence still exists for SAM with layerwise gradient normalization:
di p-n--d. d. // w Zl],
(ar, by, at di) is equivalent to (at + 0, bt 0,, ct 20,, di 01).
Under joint gradient normalization (SAM), as consider in our definition of bed- andaatereettiitiionns.. keeping the gradient norm scalings invariant under at ト a; + would require di + d; 9. but keeping the perturbation scaling invariant would require d; + di 20 as for updates. Consequently, due to the layer coupling of joint gradient normalization V| an exact equivalence between bhedd-parameterizztioos at finite width requires 0 to be the same for all layers and the conflicting gradient norm in the denominator and perturbation scaling the numerator to be accounted for by de + d; 0 and d + d - 0.
46

---

In the infinite- witth limit, abed-pan parameterizationss remain equivalent under (at + 0t. by 8,, q 28,,d, d, -2,, d) layerwise as long as the set of layers that contributes the gradient norm non- vanishingly remains invariant. The gradient norm constraints for ||u'|| O(1) become
d1 + 01 2 1/2 cv, where Vaef= = " = O(n -() with cv = min(br+1 + aL+1,CL+1 + -aa++) remains invariant under equivalence transformations.
di + 2 1 - cV.
d + a*T1 2 1/2,
Proof of Lemma F.13. First, the choices di ensure that the constraints for effective perturbations from the proof of Lemma F.11 are saturated in each layer. It is left to show, that these choices satisfy the || V L|| = @(1)-constriins. For input layers, since -d 2 a1 + 1/2, it holds that d1 + - 1/2. For hidden layers. since -d a,, it holds that di + a; 0. For output layers, since -d 2 aL++ - 1/2, it holds that di-1 + aL+1 2 1/2. Observe that the minimizer in the definition of d saturates its gradient norm constraint.
E.7 The spectral perspective jPP²
While Tensor Programs allow track the transformations of vectors like activations, Yang al. (2023a) provide an equivalent formulation in terms of weight matrix spectral norms. They find that the spectral norm measures the effect of a weight update the activations, under certain non-cancellation assumptions and limited batch size. For all MLP layers, they show that jPP is equivalent to achieving the condition
=( / fan_in fia__in
faanoou (V fin
I =
and |Aw'll- =
at all times t, where W!:: iRfan_in + _ This condition is achieved with initialization di, SGD learning rate 77 and Adam learning rate my Adam chosen as,
of =
(v Vfaa_in in
min 1.V
fan_out
fan_ in )).
7 -
fan_out
Adom-
fan_in
fan fan_ini
This generalizes μP to varying widths inside the network. For varying widths, adopt the notation W 12a4-1 - Rn with no = dn and πg+11 = dout, whereas fan_ in and fan_out always adapt to the weight matrix under consideration.
To understand why the spectral norm is desirable, note that AW m V C(x is low rank and aligned with the incoming activations. For batch size 1, we even have rank-1 updates with JAW'| | m V [1|2||=/// 112, aligned with the incoming activations =/-1 hence ||
| AAW || 1 12- This allows to achieve Ar l2 @(vmi) irrespective of the layer type with
= 00(ll==0((//m/m111
Our simple condition that perturbations should scale like updates, which is rigorously justified by our Tensor Program based proof in Appendix now allows derive the correct perturbation scalings using the spectral weight perspective.
Layerwise perturbations As simple starting point, consider a variant of (SAM) that does not globally normalize the gradient all layers jointly, but uses layerwise normalization (LN), resulting in the layerwise perturbation rule,
E'= V Lw w LII.
where ||:|| may denote either the spectral or the Frobenius norm (equivalent under limited perturbation batch size). Without the global normalization, the scalings of all layers not coupled, and the immediately requires choosing
spectral condition ||e||--
O fan_ Tan_out/fant
Pi = p. /fan_out/faa__ii
for effective perturbations in layer with width-independent hyperparameter p 2
Perturbations with global gradient normalization. Perturbations that are globally normalized across layers have usually been implemented practice according to the GitHub repositories provided
47

---

by Foret et al. (2021); Samuel (2022); Kwon et al. (2021); Andriushchenko and Flammarion (2022); Müller et al. (2024). Since we are interested in analysing (SAM) as is applied in practice, we study variants with joint gradient normalization in more detail. Preliminary ablations in Appendix H.5 suggest that layer-coupled SAM with global normalization slightly outperforms SAM with layerwise gradient normalization. To simplify the analysis as much as possible, we will first ensure width- independence of the normalization, so that the layerwise perturbation scaling is not affected by the normalization term. Then, layerwise perturbations should again be scaled like updates.
Separate denominator scalings. If we allow scale each denominator term separately from the corresponding numerator term (DP), the perturbation radius in each layer for the numerator can be fan_out
scaled like updates, I O
Ian_in
Now, to ensure (1) in the denominator, each input and hidden-like gradient norm V Cllr, I E [L], achieves width-indeeenddenc if it scaled by biases when understanding them as weights - Rn to the input These scalings are derived in the next paragraph. The last-layer gradient norm Ww+++f/l shouldd and g as L+11
//ff/aannnnnna//a The same rule applies to
should be scaled as (M.m
1/2
If we care about the correct width-independent constants, observe that the learning rate scaling ML+1 oftaaaaot/ffa_iiiii in) induces ||W++1| ||WW +| 11 = O n2 //m)) If we wanted
to achieve AWL+1 0( (((vnL+i/n we would need 7 @(1/fan_in) in). As nL+1 dout is width-independent, /faan_ Tifaa_out/ffeet out/fann_ia in would result the same width-dependent scaling for the last layer, but ignoring large constants can introduce significant distortion. For example in ImageNet K, n1+1 is large. By tuning input, hidden and output multipliers such constant distortions may be corrected. The multiplier used in the mup-package does not correct this distortion. Using a base width at which SP is recovered may also cement such spectral distortions, if no multipliers are tuned.
spectral
Derivation of gradient norm / r scalings. In this paragraph. ||. || may denote the Frobenius or spectral norm. As all matrices are of limited rank, both norms scale equivalently. As first step, [| O Vni can be reconstructed from
@(Vni/mi-i) = |AW* ||. = mill V -1111 112 n4 11/ 1 /n (l
V N: C1 -e-_|ll=~-1 O @(v m/- /n). Multiplying by scaling. The output layer gradient V V L E R=241x01 is given by (VWL+1 C)1; = = = (()), so that ||V (()TRT1))) Biases before the last layer follow the scheme V| || = V (I/mm) = @(Vfan_in caan_inffa__uute The last layer bias [ ClI = /ML+1 scales indhh-inndeendently as should, but needs to be scaled by a different constant 1/ /fan__ouuu out than earlier layers.
Now, for input and hidden layers, V W |
the inverse yields
Extensions to ASAM. As ASAM cannot be written as NESOR T program, its scaling can only be derived heuristically. As provided in Table and derived in Appendix F.4, elementwise ASAM scales all layer types correctly in relation to each other, and it suffices to rescale the global perturbation radius by /nL. assuming all width dimensions scale proportionally. For SAM-ON, we only perturb input-like layers such as normalization layers. As the conditions for correct scaling remain the same, the above scalings for input layers in SAM also apply to SAM-ON.
For layerwise ASAM, first note that ||W | = @(IlW]] = ((vn)) for input and hidden lay- ers 1 E [T]. As the numerator contains || it requires the layerwise perturbation scaling fan In the denominator, width independence is achieved with the multiplier Fan_in since
||W* |F V CIl V Ri-1 VI- Again, the output layer requires a special treatment.
Due to its small initialization, it holds that |W I| = || 1 ||3 = e (2
For perturbations
my
that fulfill the spectral condition PL+1||W +++1 || V w LlI = @(v
mLL1 my
we need to choose
PL+1 p (width-independen) but very small). The last-layer denominator term scales as
|| 1 = =e(v
"L nT.
(TLLL+1)
O(n2++)) which is width independent, but
48

---

can be large constant, as for ImageNetIK. The output bias numerator exactly conforms with the correct scaling V n1++1 fan_out faanttfan_iiar in (1))
Note that weight decay may break statements like || = ((1lW))l)) over long training. Everett et al. (2024) have recently observed more generally that scalings may evolve differently over long training than predicted by pure infinite-width TP theory, because alignments evolve dynamically between CLT- and LLN-like behaviour.
Using the mup-package. The mup-par sap-packae introduces the output layer weight multiplier ni so that input and output layer learning rates may be scaled by the same width-dependent factor. Hence, only the last-layer scalings change. The scalings of n WW+1 and nt AW remain the unique ones that achieve J.P, but V W L is scaled by "z This requires adapting the last-layer learning rate IL++1 to scale like input layers. For SAM, the last-layer perturbation radius can now be scaled like input layers. That is, assuming proportionally growing width n, in the numerator Pt.+1 = P1 = p. and pt = for E [2, LJ. and the gradient norm contributions should be scaled by vn for input and output layers, and by 1 for hidden layers. The Tensor Program perspective weight multipliers in Appendix F.6. The correct width-indeeendent constants are achieved with the last-layer numerator scaling PI-1 p. and the last-layer denominator scaling /nL/ since Wwzsı = @ 0(/nL+1//L) and for the numerator get additional nz in the forward pass. For SAM-ON nothing changes, only input-like layers are perturbed. For elementwise ASAM, ignoring width-independent constants, nothing changes the weight multiplier ny increases the weight scaling WL+1 and decreases the gradient scaling V by the same amount. The additional n -fctor in the numerator is cancelled out by the additional WL+1-factor. For the correct width- antth-nndependent constants with decoupled numerator and denominator scaling, we would scale the denominator by / /nL/n/.+1 with or without weight multiplier, and scale the numerator by
can be
PL+1 = P- n1/n22+1 with or without weight multiplier. For the example of layerwise ASAM, we Syre+fll, - (n(+11) again because the weights WL+1 are scaled up by n.. and the gradient is scaled down by the same amount. In the numerator, the upscaling of the weights also cancels out the downscaling of the gradient and additional nz in the subsequent forward pass, leading to an unchanged PL+1 potentially leads to numerical issues.
still get for the denominator [
p - 1. which is width-independent but 3
Code for yd? with separate denominator scalings. Algorithm provides a Py Torch code example that implements the above ju? scalings for SAM, scaling the gradient norm contributions of all layers to (1) (equivalent to (a-uP²2 together with naive perturbation and learning rate scaling). We adapt the popular SAM implementation Samuel (2022) using the mup-package. This code resembles our implementation for the ViT experiments. In the mup-package, 'vector-like" parameters scale as n x constant or constant X n and include input and output weights. The last-layer multiplier n -1 is chosen that input and output layers can be scaled by the same width-dependent factor. On the other hand, "matrix-like" parameters scale as n X n and include hidden weights. The implementation uses base width at which μP? and SP are equivalent; all width metth-deeendent scalings then scale with width-mmltipliers width/bass bbap//iadhhhbhh...eaeap width. This allows to immediately transfer well-performing settings from SP to pP22.
Let us recapitulate how the μP² scaling in the following code arises. The crucial variables track are factor, group [" "J and group g'gaadnomm. grranorm_scaling"" scaling"], For limited batch size, the spectral and Frobenius norm of gradients scale equivalently, and we get, for all I E [Z.].
((V
f in (Vran_
V Cw V(IV =
We want scale each weight's contribution in the denominator to be width-independent, hence need the factor / /facorr with factor fan_ dan_out/ae /////anninnn in. For the numerator, the spectral condition (*) demands ||pt Cll 1 0(, [an Ian_oot ), so that we need to scale the weight's perturbation radius fan_in
to p: = p factor. Since the mup- mup-package sets the last-layer weight multiplier such that input and output layers can be scaled in the same way, the implementation is short. For optimal numerical properties however, this choice of multipliers is sub-optimal (Blake et al., 2024).
49

---

import math, torch from mup import MuAdamW
specify parameterization
parameterization 'mupp' # 'sp-naive', 'mup- naive'
for 'mup- global use "mup-nad maieenaive' and scale rho accordingly
specify model and hyperparametere model, 1r, rho, weight_decay, wew._llyrrwweight... name
10
12
13 adapt SAN to allow gradient norm scaling each weight tensor
14 class StM(torch www.:oppttiii........... Optimizer):
17
def grad_norm (self): grads []
18
19
for 1, group in (sueraat((eeff par.....aaa..gusss)....
for in group ("pramm"
21
grads append ( (groop["gradnorn scaling"] p p.grad) norm(p-2) (grad).nnn((p1))) norm forch.stace (grads). norm(p=2)
22
23
return norm
mhorch.no_grae
26
def first_step(selff: perturbation step before the weight update grad_norm self grad_norm()
27
28
for group in self saff.ppram_ggouus
scale group["rho"] (grad_norm 1e-12) for p in group["parans"]
30
31
if p. p.grad is None: continue
32
self.state (olddpd_pp = p. data. clone
e_w p.grad to))
35
p. (e_dd climb to local maximum e(w)
36
37
width -idth-dependent rho and gradient norm scaling for each weight 40 param-groups
41
for name, in model. named_ paremeterra ():
42
if p. pnishhape ninf () == 0 or 'naive' in paramoterization:
factor =1
44
elif p.infshape. ninf ()
45
vector -lke
16
for d in p. infshaae
if d. babas__dim is not None factor d.dim / d. base___i_s #width break
49
50
elif p. infnfhapp.... ninf () 2:
# matrix-like
factor (p. infshape [0] dim/p. infshhae [1] dim)
(p. infshape [1] base base_dim/
p. infshaape [0] dim) fan_ www.ounnnaaaaaaoo
53
else: raise NotimpleemntedEFror
54
group "params": [p], "Ir": lr.
37
58
59
"rho": rho factor,
"gradnorm_scaling" math. sath.sgrt(factor)
}
62
param_groups append (group)
63
50

---

optimizer = SAM (paraan.ggroopp
66 base_ aame_oopimiizerkuuddan' if parameterizationn= mmpp else torch optim
Adamw, weight_ herpaaadddaaaseiigtdeeeeeeee
Algorithm 1: Pytorch implementation of μPP² for SAM using the mup-mackage package. Key changes from the original implementation that correct the layerwise perturbation scaling are highlighted with gray boxes. This code decouples the scalings of numerator and denominator terms following (DP), and scales the gradient norm contributions of all layers by group ("grddnorm, sarrrmssaalll scaling"] in the denominator be audth-independentt The numerator terms group ["rho"] of all weight tensors are scaled to achieve effective perturbations. This scaling is equivalent to (a-uP2) together with naive perturbation and learning rate scaling.
G Experimental details
If not mentioned otherwise, experiments use the settings specified in this section.
Implementation details. For MLPs, we exactly implement our Definition 4 of pod-parrameteizzations precisely validate our theoretical results. For ResNets and ViTs, the width varies inside the network, so that we implement the spectral scaling rules derived in Appendix F.7. Like the mup-package, we introduce base width at which SP and j(P are equivalent, allowing immediately transfer setups that perform well in SP. We use the mup-package only for ViTs, and our implementation of μp² resembles the pseudocode provided in Algorithm 1. For ResNets, we use no width-dependent last- layer multiplier. At initialization, jP differs from SP only through smaller last layer initialization. For MLPs we exactly implement the bed redeparramttrizaaionn with ب = 1, but use the large width- independent input layer initialization variance 2 instead of the width- indhh-independent 2/din. in μpP, which can be seen as tuned initialization variance multiplier. For ResNets and Vits, we initialize the last layer to 0 in JP, which corresponds to b + oo and which recovers the limit behaviour So - 0 already at finite width. We are working on making Python code to reproduce all of our experiments publicly available.
MLPs. We train 3-layer MLPs without biases with ReLU activation function for 20 epochs with constant learning rate, using SGD as base optimizer as specified in Definition 4, but allow for SGD batchsize larger than 1, defaulting to batch size 64. We evaluate the test accuracy after every epoch and use the snapshot across training with the best accuracy. This is necessary as the test accuracy is not monotonically increasing across training, while the training accuracy is. For ResNets we do not observe such harmful overfitting. For the standard parametrization, we use He initialization (He et al., 2015) and don't tune multipliers mimic standard training procedures. For j.P we resort to the optimal multipliers from Yang et al. (2022). We then find the optimal learning rate and perturbation radius for each bed paramettriiziiion and SAM variant separately.
ResNets. For 8 experiments, we augment the CIFARIO data with random crops and random horizontal flips, set labelsmoothing to 0.1 and use a cosine learning rate schedule. ResNets in uP have base width 0.5, gradient norm scaling according to Definition 4 and their last layer is initialized to O. For SP, we again adopt the standard hyperparameters from Müller et al. (2024) by using a momentum of 0.9, weight decay 0.0005, an output multiplier of 1.0, and individually tuned learning rate and perturbation radius for each SAM variant. For j.P, at base width multiplier 0.5 compared to the original width, for each SAM variant, we perform a random grid search over the hyperparameters learning rate, perturbation radius, output multiplier [2 8, 8,2 ......2. ,?%, weight decay [o, 10 5 10 4, 5- 10 10 3, 10 2] and momentum [0, 0.1, 0.4, 0.7, 0.9]. Learning rate and 2-10 2-2--999. ...... 2 centered around recommendations from the literature. The optimal hyperparameter configurations found from at least 1 150 runs for each SAM variant are summarized in Table G.I. Learning rates and perturbation radii were further tuned with the experiments from Appendix H.3.3.
perturbation radius grids were either
VITs. We train VIT-S/16 with 6 layers and 12 attention heads on ImageNet1K K (Deng et al., 2009) and a ViT-S/4 with 12 layers and 12 attention heads on CIFAR100 (Krizhevsky al., 2009) (see Appendix H.6), again adopting the hyperparameter settings from Müller et al. (2024). This means we use AdamW as base optimizer with warmup and a cosine learning rate decay. For CIFARIO0 we use random crops, random horizontal flips and AutoAugment as data augmentations. For Imagenet we use the original preprocessing from Huggingface has--bas--pccchh 222 (Wu et al., 2020). For μP,, we tune multipliers at a basewidth 384, initialize the last layer and query weights to 0. By
51

---

using the juP package, the relative perturbation scalings change explained in Appendix F.7 and Appendix F.6. Global and naive perturbation scaling in μ&P now coincide. Here, instead of the original perturbation scaling Definition 4, we scale the gradient norm contributions of all layers in the denominator to (1). The hyperparameter choices for ViTs on CIFARIO0 and ImageNet are summarized in Table G.2. For J.P, the learning rate, perturbation radius, input multiplier, output multiplier and weight decay were tuned using independent runs of Nevergrad NGOpt with budget 56 on ImageNet. The same multipliers are used on CIFAR 100.
Figures. Whenever multiple runs with independent random seeds used for training. confidence bands cover the interval from the empirical 2.5% to the empirical 97.5 57.5%-quantile. The line then denotes the average of all runs. When confidence bands given, but the number of independent runs is not specified, the number runs defaults to 4.
Computational resources. We ran all of our experiments on Amazon EC2 G5 instances each containing up to 8 NVIDIA A10G GPUs. On single GPU, our ppP2-SA training script for MLPs of width 4096 CIFARI0 10 takes 502 seconds to run in total (25 seconds per epoch), where data handling takes most of the time. The training times for ResNets and ViTs are presented Table G.3.
Hyperparam.
SAM
SAM-ON
ResNet18 SGD
Elem. ASAM
Layer ASAM SP
SP
μP2
SP
puP2
SP
pup
μp2
Training epochs
200
Batch size
64
LR 7
0.05
2-4
0.05
2-4
0.05
2-4
0.1
2-4
LR decay
Cosine 0.0005
Weight decay
Momentum
0.9
Labelsmoothing
0.1
Pert. radius
0.1 2
0.5
5-2
2
10 2-4 0.02
2
Output multiplier
0.125
0.125
0.125
0.125
Table G.I: (ResNet-18 hyperparameters for CIFARIÓ) Hyperparameters for SP are taken from Müller et al. (2024). Learning rate and perturbation radius are tuned using the experiments in Appendix H.3.3. ResNets in yP have base width 0.5, gradient norm scaling according to Definition and their last layer is initialized to O.
Hyperparam.
SAM on ImageNetIK SP
SAM on CIFAR100
nP²
shared
SP
μP2
Training epochs
100
300
Batch size
128
LR 77
0.001
0.00226
0,0005
LR warmup epochs
10
LR decay
Cosine
Weight decay
0.1
0.0872
0.05
Labelsmoothing
0.1
Pert. radius p
1
1.1939
0.25
0.25
Input multiplier
1
1.7309
1.7309
Output multiplier
4.0946
4.0946
Layers
6
12
Attention heads
12
Patch size
16
Table G.2: (Vision Transformer hyperparameters) Hyperparameters for SP are taken from Müller al. (2024) using AdamW as base optimizer. VIT's in jP have base width 384, last layer and query weights are initialized to O and gradient norm contributions of all layers are scaled to (
52

---

H
Supplemental experiments
This section provides more extensive empirical evaluations to validate the claims of the main paper. By naive perturbation scaling (naive) we denote parameterizations that do not adapt any perturbation scalings (d = di = O for all (). Global perturbation scaling (global) denotes the maximal stable scaling of the global perturbation radius that achieves effective perturbations in some layers without layerwise perturbation scaling (d, for all l).
SAM is approximately LL-SAM in jPP with global perturbation scaling H.1
Figure H.1 compares SAM in μP under global perturbation scaling (uP-global) with SAM under global perturbation scaling where only the last-layer weights are perturbed (LL-SAM) by showing more neural network statistics that are related to SAM's inductive bias and to learning in general. From top-left to bottom right, the statistics are: Frobenius norm of the layerwise weight perturbation (which closely related spectral norm as perturbations are low rank); Frobenius norm of the layerwise weight perturbation normalized by the weight spectral norm to upper bound the influence of the perturbations on the output; spectral norm of the weight updates across training scaled by the spectral condition n²/22, 1 and 11/2 for input, hidden and output layers respectively; norm of the activation updates for each layer normalized by the square root of the layer's output dimension to measure coordinatewise update scaling; layerwise effective feature ranks measured as in Andriushchenko et al. (2023a) by the minimal amount of singular values to make up 99% of the variance of the activations in a given layer; gradient norm, Hessian spectral norm and Hessian trace of loss with respect to weights; training accuracy, test accuracy after optimally stopping.
Observe that, especially for large widths, global perturbation scaling effectively only perturbs the last layer, as predicted by Theorem 11. Last-layer SAM is more similar to juP-global SAM than SGD on all of the tracked statistics, in particular large widths. Only perturbing the last layer still affects the gradients earlier layers so that weight updates and activations change in all layers. We find that SAM in μP with global scaling does not consistently improve generalization performance over SGD, whereas μu? does improve over SGD for all widths (Figure H.3). Last-layer perturbation norms coincide by design with the global perturbation radius n "pp p and their effect on the activations stays ((1) with increasing width as measured in relation to weight spectral norm. Formally the last-layer perturbation norm converges due to
| - Wh- 1|F n Pl
X Pee
+
pl
=
12 dp p-> 0,
where the loss derivative Xt always cancels out due to the normalization and the global gradient norm ||uz||| is dominated by the last-layer gradient norm due to the global scaling (Theorem 1I). Normalizing the weight perturbations by the weight spectral norm measures the influence of the perturbations on the activations. Note that this influence is also vanishing. Feature ranks stay close to initialization, since random initialization has high rank and training does low effective rank updates. Here we do not observe that SAM reduces the feature rank compared to SGD. The Hessian spectral norm and trace quite noisy. The last-layer Hessian spectral norm explodes with width in μP, because last-layer learning rate is scaled as n hence the edge of stability explodes. ResNets in jP more stable, their Hessian spectral norm even shrinks with width (not shown).
Contrast the results for jP--llobal with the results for μp² in Figure H.2 for comparison with SGD in μP. The Hessian spectral norm is reduced by SAM as you would expect. Additionally μP² shows low variability in performance and all other statistics. SAM in juP2 does not reduce the feature rank compared to SGD in pP. This suggests that the conclusions drawn by Andriushchenko et al. (2023a) do not apply MLPs pP.
ResNet-18 18 on CIFAR10
ViT on CIFAR100
ViT on ImageNet1K K
Width multiplier
0.5
2
4
0.5
2
0.5
2
Seconds per epoch
109
161
327
803
209
327
777
2550
4151
9802
Table G.3: (Training time per epoch) Training time (in seconds) per epoch of the entire data loading
and training pipeline of SAM in juP2 on a single NVIDIA A A10G GPU.
53

---

H.2 Propagating perturbations from the first layer does not inherit SAM's benefits
Here apply parametrization that only effectively perturbs the first layer weights (derived in Example F.I). Figure H.2 shows that effective first-layer SAM loses both ju2 SAM's improvement in test accuracy well SAM's inductive bias towards smaller gradient norm and Hessian norm, iee. lower sharpness MLPs. This performance deterioration occurs although the perturbation of first-layer SAM has effect of the same order of magnitude as yP2 weight and activation updates in all layers. This shows that mere propagation of weight perturbations from earlier layers cannot replace effective weight perturbations each layer order benefit from SAM. It is crucial to correctly adjust the layerwise perturbation scaling, and to distinguish between effective perturbations and perturbation nontriviality in each layer.
SAM in jP? on the other hand, achieves the correct perturbation and update scaling, has lower final gradient and Hessian spectral norm, improves test accuracy over SGD and has overall lower variance between training runs.
54

---

Ourpet
LAK15
CA0123
2.420
0.002
0.400
C0000
L1055
LA0A4
0.400
1480
1.20
02/02
1.200
100 11
3.40
3.200
C1M92
2.43
2..00
Teper kyer
Outpul
04.23
1/M0 6.60
500
Cil STELT
0.20
0.13
9180
Clapes
300
0.2
EXU 633
Width
With
344 20
143 901
.000
140
Thiningg actHinTy
100
12
25
05
60
Wideh
Whdth
Figure H.1: Several neural network statistics for SAM (blue), LL-SAM (green) and SGD as a baseline (orange) across width after training a 3-layer MLP in yP-gglobal for 20 epochs with the optimal learning rate 0.3432 and perturbation radius 0.2154. The statistics are explained in the text of Appendix H.l.
55

---

0.375
100 111
0.000
0.311
0.005
1.03
0.40
0.399
0.033
0.005
60033
6.00
0.455
With
WSB8
liper lige
1iddea bigar
Dape
11.0
12.6
101
Input liye
Ege:
Oiipee
0.33
0.35
0.30
Widb
9133 HHkke low
W1.00
1101
0535
1750
1200
300
1.50
1000
0.00
943b
Widhs
140
20
SED
5120
Thiiiing stsmmec
Figure H.2: Same neural network statistics as in Figure H.I but SAM-SGD in pP2 (blue) versus MUP with perturbations scaled to only effectively perturb the first layer weights (green) with SGD in pP as a baseline. The first-layer perturbation parameterization performs worse than ypP² and results in gradient norm and Hessian norm similar to that of SGD, larger than those of SAM. While the spectral norm of the weights converges to a similar quantity as for pP2, the effect of the weight changes on the hidden activation updates behaves more like SGD. Feature ranks all look similar.
56

---

H.3 Hyperparameter transfer
In this section, we provide supplemental evidence that, jaP² is the unique perturbation scaling that robustly achieves hyperparameter transfer in y&P both for the optimal learning rate and the optimal perturbation radius across neural architectures and datasets. But we also show that both MLPs and ResNets in SP can sometimes achieve hyperparameter transfer after long training.
H.3.1
MLPs /P
Figure H.3 shows that in pP2 the optimal hyperparameters in terms of test accuracy transfer in both learning rate and perturbation radius at sufficient width, and test accuracy monotonically improves with model scale. In addition, SAM in pP outperforms SAM in jP with global perturbation scaling all widths.
While other works focus on hyperparameter transfer in training loss, we are ultimately interested in transfer with respect to test accuracy. Especially under harmful overfitting, the test accuracy affected by nontrivial interactions between the learning rate and the perturbation radius. While the joint optimum is slightly shifting towards larger learning rate and perturbation radius for small widths, it remains remarkably stable for sufficient width 2 1024. Note that slight shifts in the optimal learning rate due to finite width biases have also been observed in earlier works (Yang et al., 2022).
100
90
Training sccuracy
Training accuracy
80
80
70
60
50
1dP
Leaming
Perturbation radius
width
58
Validatio accurac
Volidatoon ROCHEaCy
58
7034
86
4095
00.0
47.5
46.0
42.5
48
40.0
Learning
Perturbation
Figure H.3: Training accuracy (top) and test accuracy (bottom) after optimally stopping 20 epoch SAM training as function of learning rate (left) with perturbation radius p = 0.2154, and as function of perturbation radius (right) with learning rate 0.4529 in yu22 The optimal learning rate transfers. The smaller the perturbation radius the better the training accuracy. For sufficiently wide MLPs, the validation-optimal perturbation radius transfers well and SAM reduces harmful overfitting.
57

---

96
96
Triining accurec
Trinnng acturicy
88
85
75
70
60
60
00
10
10
Leurning rale
Pestrrbation radius
68
588
width
Vulidation accuraty
Validatioo accuracy
1004
4095
64
200
150
10
Learning rale
Perturlation radius
Figure H.4: Training accuracy (top) and test accuracy (bottom) after optimally stopping 20 epoch SAM training as function of learning rate (left) and perturbation radius (right) in iPP-global with the same base learning rate and perturbation radius as in Figure H.9. For global perturbation scaling, we do not observe a benefit of SAM over SGD.
90
98
Thinnng accuracy
Triining accuracy
801
88
70
75
60
65
66
66
19
Learning
radius
68
68
64
Validation accuricy
56
Validution accuracy
56
4005
00
-00
30
48
Learning
Pesurdation radius
Figure H.5: Same as Figure H.4 but with input multiplier 0.0305 and small output multiplier 0.0098. Note that networks with width at most 256 perform better in terms of test accuracy than with the other multiplier choice in Figure H.4, but the multipliers here have worse width scaling properties. To the best of our knowledge, the issue that optimally tuned hyperparameters on small models may scale worse than slightly suboptimal hyperparameters has not been stated before. This raises the question when and how can we use small models to predict the optimal hyperparameters of large models.
58

---

Figure H.4 shows that global perturbation scaling does transfer the same perturbation instability threshold, whereas in jal-naive every fixed perturbation radius becomes unstable sufficient width (Figure H.7). But in juP-global we do not observe a benefit of SAM over SGD. While the optimal learning rate with respect to the training accuracy transfers, the optimal learning rate with respect to the validation error is smaller for MLPs of moderate widths due to harmful overfitting. How to control for non-monotonic dependence of the test error on the training error is an important question for future work. Figure H.5 also shows juP-global but with different choice of input and output multipliers. With these multipliers, networks with width at most 256 perform better in terms of test accuracy than with the other multiplier choice in Figure H.4, but these multipliers have worse width scaling properties. To the best of our knowledge, the issue that optimally tuned hyperparameters on small models may scale worse than slightly suboptimal hyperparameters has not been stated before. This raises the question when and how can we use small models to predict the optimal choice of all hyperparameters jointly in large models.
L -naive
10a
Learning rate "
X
X
10
10-2
10-1
Perturbation radius p
Figure H.6: Same Figure 1 but for uP with naive width-independent perturbation scaling p. The regime of stable perturbation radii shrinks with increasing width as predicted by Proposition 1.
Figure H.7: Mean (over runs) of training accuracy (top) and of test accuracy (bottom) after optimally stopping 20 epoch SAM training of MLP in yaP-naive as a function of learning rate and perturbation radius. The optimal hyperparameters do not transfer. Every fixed perturbation radius becomes unstable in sufficiently wide networks.
59

---

Figure H.8: Mean (over 3 runs) of training accuracy (top) and of test accuracy (bottom) after optimally stopping 20 epoch SAM training of a MLP in pPP-llobal as a function of learning rate and perturbation radius. The global scaling of the perturbation radius by n -1/2 compared to yP-naive (Figure H.7) makes the stable regime invariant to width. But the suboptimal layerwise perturbation scaling that only perturbs the last layer does not consistently improve over SGD (p- 0).
Figure H.9: Mean (over 3 runs) of training accuracy (top) and of test accuracy (bottom) after optimally stopping 20 epoch SAM training of a MLP in up² as a function of learning rate and perturbation radius. At sufficient width, the optimal hyperparameters are stable in terms of test accuracy, even under severe overfitting.
H.3.2 Some variants of SP can transfer optimal hyperparameters on CIFAR-10
Surprisingly, after training MLPs to convergence on CIFAR-10, some variants of SP with naive perturbation scaling can transfer learning rate and perturbation radius, against the prediction by infinite-wind width theory. For SGD, this has originally been observed in GitHub issue 52 of the mup- package*. We train MLPs with SAM in variants of SP in Figure 10 and also observe that some variants achieve transfer while for others the optimal learning rate shrinks as in Yang et al. (2022).
To achieve shrinking stable and optimal learning rates in SP. Yang al. (2022) weight multipliers tuned at base width 256 and normalize the initialization variance to be invariant to these weight multipliers, according to the Jupyter notebook provided for reproducing their experiments. In addition, they initialize the last layer to 0 which contradiets SP scaling but results in more striking shrinkage of the optimal learning rate. We observe that both with and without weight multipliers, MLPs trained with SAM in SP-naive have surprisingly good transfer properties on CIFAR-IO. 10. With tuned multipliers but initialization that is invariant these multipliers, the optimal learning rate shrinks. Because we are training to convergence, pure infinite-width theory does not adequately describe the training dynamics anymore (Vyas et al., 2024). Infinite-width theory implies that scaling the width further would eventually break the learning rate transfer. It remains matter of ongoing work to understand whether this stability of SP is a finite-width or long training time effect, and whether this empirical stability is particular to multi-epoch training on vision datasets. As shrinkage of the optimal learning rate in SP has generally been observed in language settings (see e.g. Brown et al., 2020, Table 2),), we expect the same shrinkage for SAM in SP in such settings.
'Without tuned weight multipliers. MLPs trained with SGD SP on CIFAR-10 can transfer the optimal learning rate: https: httpt::///gttttbb wwtp://oo//ooooo///o/oooooooo wwtp///roooottttiiooe/eeeeeeeeee 'https:/ /ttpgggtttttttt////e/eeeeeeehh wwtp////ccroooooo//////oo///abaa mawi//////mm/////aeeaa/aaaaaaeaa www.peeeeeeeeee///eeeeeeeeeeeeee MM////LPPP/ ipynb
60

---

Note that the learning rate transfer in SP here is a much stronger observation than in Everett et al. (2024) who choose the correct layerwise learning rates for SP. Hence their SP merely deviates from juP through larger output layer initialization and key-query normalization by v SP versus by d in μPP Here however we even observe transfer in our stricter understanding of SP without any layerwise learning rates or weight multipliers. This is not peculiarity SAM; we observe the same learning rate transfer in plain SGD without any momentum or weight decay for MLPs CIFAR-10 (not shown).
SP-naive. mults
SP-naive (Yang mults)
SP-naive, initnorm
SP-naive. initnorm, Il=0
Leaning nite
witth 0396
1120
radius
Figure H.10: Optimal learning rate and perturbation radius (cross) and regions within 1% of the optimal test accuracy (mean over runs) after optimally stopping 20 epoch SAM training of MLP in different variants of SP for varying widths (the darker, the wider). Observe transfer properties in SP almost stable in μP2 (Figure 1) and against infinite-width predictions slightly growing with width, both with and without the tuned weight multipliers by Yang et al. (2022). Only when normalizing the initialization variance be independent of the width-independent weight multipliers (initnorm), does the regime of stable learning rates shrink at the widths considered. Additionally initializing the last layer to zero (Il= 0) (as in the Jupyter notebook provided to reproduce Figure 3 in Yang et al. (2022)) shows even more pronounced learning rate shrinkage, but does not correspond to SP scaling anymore.
ResNets SP show hyperparameter transfer across most SAM variants too, as soon as we tune momentum, weight decay and labelsmoothing (Figure H.11). This is in line with previous empirical observations (Yang et al., 2022, Figure 16 for SGD) but contradicts infinite-width theory as for MLPs.
London
(a) momentum, weight decay or labelsmoothing
(b) Tuned momentum and weight decay
Figure H.11: Training accuracy (top) and test accuracy (bottom) after optimally stopping 100 epoch SAM training as a function of learning rate and perturbation radius in SP-naive without regularization (left) and with tuned regularization (right) using momentum 0.9. weightdecay 0.0005 and labelsmoothing 0. CI denote the minimal and maximal value from 4 independent runs. Without regularization, the optimal learning rate shrinks with width. Given the learning rate, the optimal perturbation radius seems quite stable, but since the optimal learning rate shifts, the performance scales worse than for μP² with the fixed learning rate that is tuned on the small model. With optimal regularization, both optimal learning rate and perturbation radius remain remarkably stable. We plan investigate this mechanism an upcoming work.
H.3.3 ResNets
In this section, we plot averages and a-CI from 2 independent runs.
61

---

ResNets in pp2 transfer both the optimal learning rate and perturbation radius for SAM (Figure H.) SAM-ON (Figure H.14) and elementwise ASAM (Figure H.15), well as different alternatives of scaling the gradient norm contributions to SAM's denominator (Figure H.18). This suggests correctness of the derived scalings. width multipliers 2 and 4, pP2 achieves the same or slightly better test accuracy than SP in all SAM variants.
Figure H.13 shows ResNets trained with SAM in different parameterizations. In ResNets of practical scale, p remains quite stable in juP2 but surprisingly also in SP-NAIVE. In jdP, for naive perturbation scaling the regime of stable perturbation radii shrinks, for global perturbation scaling, the optimal perturbation radius shifts, approaching its maximal stable value, which stays invariant to width scaling. Here, it would be interesting to whether even larger width would lead suboptimal performance of pPP-global. ju2 is most robust to the choice of p and achieves the best test accuracy.
Leaing
40%
(a) SP, no scaling
(b) pPp
Figure H21 Training accuracy (top) and test accuracy (bottom) after optimally stopping 200 epoch SAM training as function of learning rate and of perturbation radius in SP (left) and in juP2 (right) with optimized momentum 0.9, weight decay 10-4 44 and labelsmoothing 0.1 for both μP2 and SP. In uP= the base learning rate is 1 2 and the base perturbation radius is and p = 0.1, respectively. Observe monotonic improvement with width in both training and test error. Optimal hyperparameters transfer across widths, surprisingly in both uP? and SP.
2 in SP 0.05
SP-naive
J P-nuive
j.P P-gllbbll
nP2
Petubattin y
Petuitatio eadiss
Postubatio cadias
Petuluzion cadias
Figure H.13: Test accuracy after optimally stopping 200 epoch SAM training as function of perturbation radius in various parameterizations. Dashed lines denote the base optimizer SGD with tuned momentum and weight decay the respective parameterization.
H.3.4 ASAM variants
As we are not aware of any use of ASAM with MLPs in the literature and since the amount of necessary experiments for ViTs exceeds our computational budget, we only show that ResNets trained with the all of the discussed SAM variants in pP² transfer the optimal (n, p).
For the examples of elementwise ASAM and SAM-ON the global perturbation scaling 1/2 suffices reach μP2 The stability of the optimal perturbation radius in the applied scaling 1/2 shows that in μP with naive perturbation scaling the optimal perturbation radius would grow as n 1/2.
See the previous section, for discussion of the remarkable stability of ResNets in SP. For the example of clementwise ASAM in SP, the optimal perturbation radius seems to grow.
For layerwise ASAM (Figure H.1 the optimal perturbation radius seems to grow in both SP and μP? suggesting that our scaling condition does not perfectly apply to this variant, although uP2 (00.83)) still outperforms SP (96.8640.05 (+0.83)) in terms of the optimal test
62

---

accuracy. As Frobenius norms of weights are the only component that is not representable NE@OR program, these Frobenius norms appear to scale differently than heuristically predicted over the course of training.
2011
(a) SP. no p scaling
(b) pPp
Figure H.14: Same as Figure H.12 but for SAM-ON in SP without perturbation scaling (left) and in uP (right). Both optimal learning rate and perturbation radius are remarkably stable in both μP² and SP. Since up² for SAM-ON is just μP with global perturbation scaling 1/2, transfer here implics that μP with width-independent scaling would not transfer.
(a) SP, Do P scaling
(b) pP
Figure H.15: Same as Figure H.12 but elementwise ASAM in SP without perturbation scaling (left) and in yP² (right). Observe a consistent HP landscape in uP2 but growing optimal perturbation radius SP without perturbation scaling.
2017
(a) SP. no P scaling
Figure H.16: Same as Figure H.12 but for layerwise ASAM in SP without perturbation scaling (left) and in 2P² (right). For layerwise ASAM, both μP2 and SP seem transfer the optimal learning rate as well as perturbation radius.
63

---

H.4 Gradient norm contributions have negligible effects on generalization performance
In this section we provide ablations concerning the question which layers should contribute non- vanishingly to the gradient norm in the denominator of the layerwise SAM perturbation rule (LP).
For MLPs, in Figure H.17 we scale all contributions to (1), and then set the contribution of individual layers zero, one by one. We observe no significant effect on the optimal test loss or hyperparameter transfer for MLPs. Any layer's contribution to the gradient normalization in the denominator of the SAM update rule can be set to 0 without significant effect on the test loss. This raises the question which effect the gradient normalization has in μP. Does it contribute scaling correction SP, but may be dropped entirely in jPP?
Scaling the gradient norm contributions of all layers to @(1) (first row) and then setting the first layer gradient norm to 0 (2nd row), respectively the hidden layer (3rd row), last-layer (4th row). Each individual layer seems to have vanishing contribution to the optimal test error.
Figure 1.17:
For ResNets, Figure H.18(a) shows accuracies when rescaling all layers' gradient norms to (()) and Figure H.18(b) shows the results when using the original global gradient norm rescaled to (11). Again, both methods achieve similar optimal test accuracy. The first variant shows cleaner hyperparameter transfer and monotonous improvement with width. When comparing to our original definition (LP) in Figure H. optimal performance is similar but rescaling all layers' gradient norm contributions to ((1) may even produce a slightly more stable hyperparameter-lonn landscape for ResNets.
64

---

A011
(a) Rescaling of SAM's denominator terms to ((1)
(b) Global
V L||
scaling
Figure H.18: Same Figure H.12 but with scaling of the gradient norms in the SAM perturbation (LP) denominator that scales all terms ((1) (left) and only global denominator scaling [| (right). All denominator scalings achieve similar optimal accuracy, show HP transfer in learning rate and monotonic test accuracy improvement with width. In global denominator scaling, the optimal shifts with width.
"
H.S SAM with layerwise gradient normalization
Here we consider SAM without the gradient normalization over all layers jointly. Instead we apply the layerwise perturbation rule presented in Appendix F.7,
cf (() www.
In SP, we consider a global constant Pt = P. whereas μP² the spectral condition (+) requires
Pt = p / Ffan_out/fan_in
Overall, SAM without layer coupling performs decently, but is outperformed by the original SAM in particular ResNets, in yp² and large width. But note that for ResNets adopt the hyperparam- eters tuned for the original SAM with layer coupling, so that these ablations only serve as preliminary experiments.
MIPs. SAM without layer coupling achieves similar optimal generalization in pP2 at each width compared to Figure H.9. The regime of stable (1) p) stays width-independent, but does not transfer the optimum consistently. This suggests complex or noisy dependence of the training dynamics on p.
(a) μpp
Figure H.19: (SAM with layerwise normalization MLPs) Test accuracy as a function of learning rate 77 and perturbation radius p for optimally-stopped MLP trained with SAM with layerwise normalization.
ResNets. Figure H.20 shows that decoupled SAM has decent performance, but is worse than original SAM with global normalization (Figure H.13) in both and μP22 particular at large width. As expected, o transfers in uP².
65

---

Test accuancc
06
96
96
Test
4.0
10
10-2
10
10
Porturbatinn radios
Peturbation radios
(a)
(b) uP2
Figure H.20: (SAM with layerwise normalization in ResNets) Test accuracy function of perturbation radius for ResNets trained with SAM with layerwise normalization.
H.6 Test error over the course of training
Figure H.21 shows the test error of ResNets and VITs over the course of training. jP2 always achieves the best final test accuracy. ResNets it also achieves a decent test accuracy the fastest and removes training instabilities of SAM in SP. While SGD in μP alone cannot compete with SAM in SP, SAM in jP? uniformly dominates over the entire course of training. Our theory suggests that in pP² the gradients scaled correctly from the beginning, whereas in SP they have to self-stabilize first, which slows down convergence. We plan a closer analysis in an upcoming work.
In VīTs, JuP generally achieves decent accuracy faster than SP, since gradient norms are already scaled correctly initialization. SAM converges slower than the base optimizer AdamW in favor of drifting towards a better generalizing local minimum or saddle point. For ViTs at this moderate scale, SAM in SP catches up SAM in pP² at the end of training.
10a
SAM μP2
SAM μP2
100% accurscy
SGD jaP
100% 10st accuracy
AdamW JP
SAM SP
8 10
SAM SP
SGD SP
AdamW SP
301
10
75
Fpochs of Training
Egochs of training
(a) ResNet on CIFAR10
(b) ViT on CIFAR100
Figure H.21: Training ResNet- 18 with width multiplier 2 on CIFARIO (left) and VIT with width multiplier 2 on CIFAR100 (right). SGD and Adam W are the respective base optimizers.
66

---

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: In the abstract and introduction we state our main contributions while ac- knowledging related work. All main claims are theoretically proven and/or empirically verified.
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 the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as 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: Limitations discussed in the future work section well as in the section in the appendix that is related to the respective limitation.
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 any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-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 on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably 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, a 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 impor- tant 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]
Justification: We state all assumptions in the main paper and Appendix C, and provide all formal proofs in Appendix E.
Guidelines:
The answer NA means that the paper does not include theoretical results.
67

---

All the theorems, formulas, and proofs in the paper should be numbered and cross
referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in 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 or 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 of 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 not)?
Answer: [Yes]
Justification: All experimental details are disclosed in Appendix G. Our perturbation scaling rules are clearly stated in the main paper. Their implementation with flexible fan_ in and fan_out is explained in Appendix F.7, together with pseudocode for implementing our proposed scaling rule.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a 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 are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, 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 replicate the results, access hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS 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 a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily 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 way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model limited in some way (c.g., 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, described in supplemental material?
Answer: [No]
Justification: We only propose width-dependent scaling of hyperparameters of an existing optimization algorithm. This can be easily implemented by following the scaling rules that we clearly specify in the main paper. In Appendix F.7 we even provide a code example that
68

---

contains the essential modifications. We are working on making Python code to reproduce all of our experiments publicly available. Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
aublic/guides/Codee www.ooouaisiisioooooooooooooooee
for more details.
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 new open-source benchmark).
The instructions should contain the exact command and environment needed to rum to reproduce the results. See the NeurlPS code and data submission guidelines (https: / /nips... mnips.cc/pili wwtppcccccccccccccc.cccccccc..le gww.d/gg///e/eeeee/oooooo wwweeeuubbisssiopooo
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.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended 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: All experimental details are disclosed in the main paper Appendix G. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary 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 the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: As stated in Appendix G, we repeat all main experiments with multiple independent runs and report confidence bands within the empirical 2.5% 97.5%-quantiles When we repeat experiments on Vision Transformers that we have also conducted on MLPs or ResNett, we do not use multiple runs due to limitations in computational resources. 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, confi- dence intervals, or statistical significance tests, at 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 is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
69

---

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 range (e.g. negative
error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: We provide the type of GPU used and number of GPU seconds required for each experiment in Appendix G.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: http:::/herrris..ccppblll////hi
Answer: [Yes]
Justification: We provide theoretical analysis of a widely used optimization algorithm, point out the algorithm's 's limitations in large models and propose a correction. We do not foresee any ethical concerns.
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 a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to 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: This paper provides fundamental research toward understanding and improving existing optimization algorithms for neural networks. We do not release any model or data and do not consider generative models.
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 why the paper does not address societal impact.
Examples 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 be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out, For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the
70

---

technology is being used as intended but gives incorrect results, and harms following from (intentional or 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA]
Question:
Justification: We only use standard vision architectures and vision datasets in our experi- ments and do not release any data models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or 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 to access the model or 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 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 standard CIFARIO, CIFARI00 (Krizhevsky al., 2009) and ImageNetIK (Deng et al., 2009) datasets following the standard practice. We also cite the Python assets PyTorch (Paszke al., 2019), mup (Yang al., 2022) and the GitHub repository implementing SAM (Samuel, 2022) that we use as basis for our experiments. 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 dataset.
The authors should state which version of the asset is used and, if possible, include a URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.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 thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information not available online, the authors are encouraged to reach out to the asset's is creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
71

---

The paper should discuss whether and how consent was obtained from people whose asset 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 the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research 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 contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According 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: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If 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 to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
72


Reviews:
[{"review_id": "l2Nm97u3pJ", "reviewer": "Reviewer_CqBd", "summary": "Sharpness Aware Minimization (SAM) improves performance across various neural architectures and datasets, but understanding its scaling behavior as models grow is crucial. This study examines the infinite-width limit of neural networks trained with SAM using the Tensor Programs framework. Findings show that in wide neural networks, SAM's dynamics effectively reduce to applying SAM only in the last layer. The authors propose a new parameterization called maximal update and perturbation parameterization ensures effective feature learning and perturbation across all layers. Experiments with MLPs, ResNets, and Vision Transformers confirm the method's effectiveness.", "questions": "N/A", "limitations": "Yes.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper is well-written.\n- This paper offers a robust theoretical foundation, thoroughly explaining the concepts and methodologies used.", "weaknesses": "- In the experiments, the authors should consider including more SAM-variant methods, such as ESAM[1] and GSAM[2].\n\n[1] Efficient Sharpness-aware Minimization for Improved Training of Neural Networks.\n\n[2] Surrogate Gap Minimization Improves Sharpness-Aware Training.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.", "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": "mlXTOm47WF", "reviewer": "Reviewer_egBf", "summary": "The authors extend muP based learning rate transfer to the extra gradient ascent step involved in the SAM algorithm. The authors present \"tensor programs\" theory to derive this scaling, and present convincing experiments that in their method, both step sizes in SAM transfer across width (learning rate and perturbation radius). In figure 1, the author's method also seems to get better test performance than the more naive methods.", "questions": "Some minor things to fix:\n- in Figure 1, the legend item for muP-global is unclear since it doesn't look like a dashed line\n- in Figure 2, why not just plot the relative change in weights in spectral norm?", "limitations": "The authors do include some discussion of some limitations in the future work section.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- the authors did an honest and thorough effort to port tensor programs theory to the case of SAM\n- the scalings the authors worked out seem to work really well, both achieving stable step sizes across width, and achieving really good test performance\n- I like the narrative, first showing that standard training does not effectively perturb all layers, then showing how to fix this", "weaknesses": "I like the paper and think that it should probably be accepted. However, I can see weaknesses in terms of presentation and use of technical tools. I believe these weaknesses are severely limiting the user-friendliness of the paper. I believe that the current presentation will alienate 99.9% of the community, and addressing these issues would significantly strengthen the paper and make it useful to a significantly wider portion of the community. Since I feel these weaknesses are important, depending on the author response, I am willing to either lower or increase my score.\n\n### **Analysis seems over-complicated**\n\nThe authors present \"tensor programs\" theory to derive their scalings. However, it has now been shown that muP is equivalent to initializing matrices and scaling weight updates to have spectral norm proportional to sqrt(fan-out / fan-in). This feels like a dramatic simplification and is an extremely simple condition that is easy for almost everyone in the community to understand. I see that you discuss this condition in an appendix so are certainly aware of it. Could you clarify: is your paper just fixing SAM ascent steps to have spectral norm proportional to sqrt(fan-out / fan-in)? If so, I think it would be *extremely* helpful for the reader to state this clearly and concisely in the introduction of the paper. I would actually consider focusing the analysis around this condition, and relegating most or all of the tensor programs theory to the appendix.\n\n### **Misleading theoretical statements**\n\nI feel like this paper has inherited some of the overly grandiose language from the tensor programs papers. For example, referring to your parameterization as the \"unique\" one that works feels misleading to me. Actually I can give you a variety of different ways of parameterizing layers that would do the trick. Also the statement \"It is straightforward to extend our theory to any architecture that is representable as a NE⊗OR⊤program including ResNets and Transformers\" feels unhelpful. Can you either link directly to the appendix which does these \"straightforward\" extensions, or omit this sentence? By the way, if all you are doing is scaling the spectral norm of updates, then it's obviously straightforward to extend this to other layer types---but this does not need \"NE⊗OR⊤\" programs...\n\n### **Missing related work**\n\nCould you take a look at this paper: https://arxiv.org/abs/2002.03432. It deals precisely with the question of ensuring that layers are \"effectively perturbed\" as you say, and it is prior work to muP. I believe that the analytical strategies developed in that paper could help simplify your work. The slight issue with that work is that Frobenius norms are used instead of spectral norms.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n**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).\n\nHowever, 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.\n\n**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](https://sites.google.com/view/ngsmworkshop/accepted-papers) “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.\n\nThe 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.\n\nGiven 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.\n\n**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.\n\n**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.\n\n**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.\n\n**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).\n\n**References:**\n\n[1] Monzio Compagnoni et al. *An SDE for modeling SAM: Theory and insights,* 2023.\n\n[2] Alan Edelman. *Eigenvalues and condition numbers of random matrices,* 1988.\n\n[3] Yang et al. *A spectral condition for feature learning,* 2023.", "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": "zxlVpiBIsC", "reviewer": "Reviewer_Tz6c", "summary": "This paper analyzes Sharpness-Aware Minimization (SAM) in the infinite-width limit using tensor program theory. The authors identify issues with standard SAM implementations in wide networks and propose a new parameterization called μP^2 to address these problems. They provide theoretical analysis and conduct extensive experiments on MLPs, ResNets, and Vision Transformers to validate their findings. The μP^2 parameterization is shown to achieve hyperparameter transfer across model scales and improve generalization performance.", "questions": "- Recent work has provided theoretical insights into Sharpness-Aware Minimization (SAM) beyond its initial sharpness-based interpretation [1,2,3]. Within the µP2 framework, what distinguishes SAM from Stochastic Gradient Descent (SGD)? Does the new framework also offer a fresh perspective on the generalization performance of SAM?\n\n- Is µP2 capable of facilitating depth parameter transfer? [4]\n\nReferences:\n\n[1] Andriushchenko, M., et al. (2024). \"Sharpness-aware minimization leads to low-rank features.\" Advances in Neural Information Processing Systems 36.\n\n[2] Wen, Y., et al. (2024). \"Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization.\" Advances in Neural Information Processing Systems 36.\n\n[3] Chen, Z., et al. (2024). \"Why does sharpness-aware minimization generalize better than SGD?\" Advances in Neural Information Processing Systems 36.\n\n[4] Yang, G., et al. (2023). \"Tensor programs VI: Feature learning in infinite-depth neural networks.\" arXiv preprint arXiv:2310.02244.\"", "limitations": "Yes, the authors adequately addressed the limitations", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper provides a rigorous theoretical analysis of Sharpness-Aware Minimization (SAM) in the infinite-width limit using tensor program theory. This extends the community's understanding of SAM's behavior in large neural networks.\n\n- The paper identifies the degenerate issue with standard SAM implementations in in infinite neural networks and propose a new parameterization (μP^2) to address the problem.\n\n- Extensive empirical experiments are conducted to validate the theoretical findings and demonstrate improved performance of μP^2.", "weaknesses": "- The theoretical analysis extends tensor program theory; however, the authors did not clearly introduce abc-parameterization in the main body of the paper. It would be beneficial if the authors added a notation section and provided a more detailed comparison between abc-parameterization and bcd parameterization.\n\n- The batch size of SAM is not adequately discussed in the main paper. For SAM with batch size $m$, it is called $m$-SAM, which is crucial in SAM's behavior [1].\n\n- Layerwise perturbation scaling for SAM is not a novel concept.\n\nReferences:\n[1] Foret, Pierre, et al. \"Sharpness-aware minimization for efficiently improving generalization.\" arXiv preprint arXiv:2010.01412 (2020).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**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$.\n\n**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.\n\n**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.\n\n**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.\n\n**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.\n\n**Depth transfer.** While for SGD and Adam, depth transfer has been achieved in several papers with a simple $1/\\sqrt{\\texttt{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.\n\n\n**References:**\n\nGreg Yang and Edward Hu. *Feature learning in infinite width neural networks*, ICML 2021.", "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": "JwknyyJjhZ", "reviewer": "Reviewer_dLWn", "summary": "This paper proposed $\\mu P^2$ which is an effective way to scale the perturbation radius of SAM for each layer so that the optimal hyperparameter (learning rate $\\eta$ and perturbation radius $\\rho$) transfers across different widths. Experiments show that this approach indeed allows the transfer of optimal $\\eta$ and $\\rho$ across widths.", "questions": "1. Can you provide a more intuitive way to explain $\\mu P^2$? To me there must be an easier way to present this idea.\n2. Exactly how can I set the perturbation radius based on $\\mu P^2$? Can you provide an example for this?", "limitations": "The authors somewhat address the limitations in the last section.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "- The idea is original, significant and seems to work as shown by the experiments", "weaknesses": "To me the main issue is presentation:\n- The writing style is quite verbose which makes it hard to remember all the details and the main message of the paper. \n- There is a lack of figures to help reader understand the theorem and intuition of the paper. \n- The paper is written in a way which assumes that the readers are very familiar with Tensor program.\n- Some parts of the paper are quite confusing. For instance, I find myself keep rereading the part from Line 183 to 196 about non-vanishing perturbations vs effective perturbations and I still can not tell the different between the two. \n- The definitions are incomplete also. For instance, definition B.1 in the appendix include a clear definition for $\\Theta(n^{-a})$ but not $O(n^{-a})$ and $\\Omega(n^{-a})$ and I try my best to guess what they mean but fail to do so. Please give a clear and complete definition for each notation for such a math heavy paper. Also please provide a table summarizing all the notations used in the paper.\n- All the definitions are provided in the appendix (which are optional) which means readers cannot comprehend the paper without reading the appendix.\n- After reading through the paper multiple times, it is still unclear to me how I can set the perturbation radius $\\rho$ for each layer.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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’.\n\n**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).\n\n**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).\n\n**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.\n\n**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.\n\n**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:\n\n(1) explain how to set layerwise learning rates and perturbation radii to achieve $\\mu P^2$ in the main text,\n\n(2) make Pytorch code to reproduce all of our experiments publicly available upon acceptance,\n\n(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.\n\nConcretely, 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).\n\n| |Input-like|Hidden-like|Output-like|\n| :--- | :---: |  :---: | :---: |\n| Learning rate scaling factor $n^{-c_l}$|$n$|1|$n^{-1}$|\n| Perturbation scaling factor $n^{-d_l}$|$n^{1/2}$|$n^{-1/2}$|$n^{-3/2}$|\n\nNow to apply $\\mu P^2$, we will explain the following steps in the main text:\n\n1. Parameterize the network and the SAM update rule according to $\\mu P^2$ as explained above.\n2. Tune the learning rate $\\eta$ and perturbation radius $\\rho$ on a small model.\n3. Train the large model only once using the optimal learning rate and perturbation radius from the small model.\n\n**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.\n\n**References:**\n\nArthur Jacot, Franck Gabriel, and Clément Hongler. *Neural tangent kernel: Convergence and generalization in neural networks,* NeurIPS 2018.", "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 are thankful for all of the thoughtful comments and constructive feedback to improve the clarity of our paper’s presentation. We are delighted to have received overwhelmingly positive feedback about our content and results, and we will do our best to improve the accessibility of the paper — in our own interest. The main changes in this regard include:

**Spectral perspective versus TP theory.** While Tensor Program (TP) theory plays a crucial role in our proofs (see our response to Reviewer egBf for more details), we will further highlight the spectral perspective to improve clarity and be accessible to a larger audience. Specifically, we plan to make the following concrete changes:

1. We will introduce the spectral condition already in the introduction to intuitively explain effective perturbations.
2. We will append the sqrt(fan-out / fan-in) scaling condition for perturbations to the intuitive perturbation scaling condition after line 271.
3. When discussing non-vanishing versus effective perturbations, we will add the spectral weight scaling condition $||\varepsilon^l ||\_\ast/||W^l ||_\ast =\Theta(1)$ for all l.

**Open source code.** Upon acceptance, we will release open source code to reproduce all of our experiments and to provide another resource for understanding and experimenting with our layerwise perturbation scaling.

**Phase characterization figure.** We will provide another figure to visualize the regimes of stability, effective SGD dynamics, non-vanishing perturbations, and effective perturbations (see the attached pdf).

**Further changes to improve clarity.** In several places, we will rewrite paragraphs to improve their clarity. As motivated by Reviewer dLWn, we will describe how to use $\mu P^2$ in practice in the main paper; we will refer to the pseudo-code in Appendix E.8, and rewrite that appendix to more clearly state how to implement $\mu P^2$: First ensure that SAM’s denominator is scaled to be width-independent, then in the numerator the perturbations of each layer should be scaled like the layer’s updates. In the appendix, we will add a table summarizing all notation.

Meta-review: This paper theoretically demonstrates that weight perturbations induced by the SAM optimizer vanish in the infinite-width limit in every layer except the output layer. This may limit the regularization effects one gets from SAM in large-width settings. To address this, the theory further shows that perturbation scaling must be width-dependent. Based on this, the paper proposes a fix, called "Maximal Update and Perturbation Parameterization" (µP^2), and claims that it can achieve better hyperparameter transfer and improved generalization performance. The latter claims are empirically demonstrated in MLPs and ResNets on CIFAR10 and ViTs on Imagenet1K.

The paper received 4 reviews. Three of these reviews recommend acceptance (with different degrees), while one reviewer (dLWn) rates it as borderline reject. Reviewer dLWn's main concerns were around clarity of presentation, specifically using more figures to provide intuition about the theoretical results, as well as providing more background on the Tensor program, and finally a clearer presentation about how practitioners can use the results of this paper to set the perturbation radius for each layer. I find these to be excellent suggestions, and I strongly encourage the authors to consider them for their final version of the paper to improve its readability.

Reviewer egBf and the authors had a lengthy engagement during the rebuttal phase. The reviewer had a very positive opinion about the submission in the earlier stages of the discussion period. However, there was considerable further discussion on the definition of rigor used in the paper and how the authors and the reviewer interpret it, where there was some disagreement. Specifically, reviewer egBf suggested the possibility of a simpler proof, while the authors believed that a rigorous proof should be based on the tensor program (something that egBf disagreed with). While an alternate proof approach, especially one with the potential of being simpler, is of interest, it can be a matter of future work, specially because the current proof approach in not incorrect. Ultimately, reviewer egBf rates the paper as a weak accept.

The other two reviewers had fewer concerns about the paper and rated it as weak accept and accept. Overall, the consensus leans towards acceptance, and I also recommend accepting it because I believe both the problem and the proposed approach (Tensor program) are interesting. However, The paper can greatly benefit from a revision for its final version, with the goal of clearer presentation, based on the feedback provided during this review phase.
Decision: Accept (poster)