Title: $\boldsymbol{\mu}\mathbf{P^2}$: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling
Venue: NeurIPS 2024
Accepted: yes
Keywords: Deep Learning Theory

Optimal Hyperparameter Transfer

Sharpness Aware Minimization

Infinite Width Limits

Signal Propagation Theory

Tensor Programs
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 2 : Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling

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 ( µ 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 µ P 2 achieves hyperparameter transfer of the joint optimum of learning rate and perturbation radius across model scales. Moreover, we provide an intuitive condition to derive µ P 2 for other perturbation rules like Adaptive SAM and SAM-ON, also ensuring balanced perturbation effects across all layers.

1 Introduction

Sharpness Aware Minimization (SAM) (Foret et al., 2021) and its variants (Kwon et al., 2021; Müller et al., 2024) improves generalization performance across a wide range of neural architectures and datasets (Chen et al., 2021; Kaddour et al., 2022). In the SAM formulation, we minimize a given loss L between our prediction and the data y as a function of the architecture's weights W , where an adversary simultaneously maximizes the same loss by perturbing the weights within a budget ρ .

A standard SAM update for an L -hidden layer multi layer perceptron (MLP) is given by

where W = [ W 1 , . . . , W L +1 ] , t is the iteration count and ε l t denotes the perturbation in the l -th MLP layer with width n ∈ N , and where we define an L -hidden layer MLP iteratively via

for inputs ξ ∈ R d in with trainable weight matrices W 1 ∈ R n × d in , W l ∈ R n × n for l ∈ [2 , L ] , and W L +1 ∈ R d out × n . We call h l preactivations, x l activations, and f ( ξ ) output function. Despite the inherent difficulty of non-convex, non-concave optimization, SAM is quite successful in practice.

10

-

2

10

-

1

10 0

Perturbation radius

ρ

10

-

1

10 0

Learning rate

η

µ

P-global

10

-

1

10 0

Perturbation radius

ρ

10

-

1

10 0

µ

P 2

width

4096

1024

256

10

-

1

10 0

Perturbation radius

ρ

55.0

55.5

56.0

56.5

57.0

57.5

58.0

58.5

59.0

59.5

Test accuracy

µ

P 2

µ

P-global

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 in the infinite-width limit.

Instead, a different scaling of layerwise initialization variances and learning rates, termed Maximal Update Parameterization ( µ P), is necessary to achieve feature learning in all layers in wide networks. A crucial practical benefit of µ P is the transferability 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:

Weshow that training an MLP with the standard (SAM) update rule is equivalent to applying perturbations only in the last layer in 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 µ P.

We postulate that jointly transferring the optimal learning rate η and perturbation radius ρ requires width-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-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.

We demonstrate that the optimal perturbation radius can shift significantly in µ P (Figure 1).

We derive the Maximal Update and Perturbation Parameterization ( µ P 2 ) that achieves both feature learning and effective perturbations in all layers in the infinite-width limit. We empirically demonstrate that µ P 2 achieves hyperparameter transfer in both learning rate η and perturbation radius ρ (Figure 1).

We provide a 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 in a Tensor Program.

2 Background and related work

We here provide a short summary of related work. A more detailed account is provided in Appendix B.

Sharpness Aware Minimization. SAMwasmotivated as an inductive bias towards flatter minima and it provably reduces properties of the Hessian that are related to sharpness in simpler settings (Bartlett et al., 2023; Wen et al., 2023; Monzio Compagnoni et al., 2023). However a 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 al. (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 architectures, optimization algorithms and arbitrary abc -parameterizations. Each abc -parameterization is essentially defined by a layerwise scaling of initialization variance and learning rate as a function of network width. Beyond pure infinite-width limits, the simple 1 √ L -scaling allows depth-scaling in ResNets and unlocks hyperparameter transfer across depths (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.

Look-LayerSAM (Liu et al., 2022) already considers layerwise perturbation scaling with the goal of preserving good performance under large batch training. However, achieving µ P 2 with LookLayerSAM 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 our results for MLPs under µ P:

By analyzing the infinite-width behaviour of the SAM update rule, we show that the training dynamics under standard (SAM) become unstable as the network width increases. This result is first stated informally below in Proposition 1 and then more formally in the next section.

Proposition 1 ( Instability of standard SAM parameterization in wide neural networks ) . Under µ P 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-width limit for any fixed, positive learning rate η > 0 and perturbation radius ρ > 0 .

Hence, to achieve stable optimization, it is necessary to introduce some width-dependent perturbation scaling ρn -d for some suitable d > 0 . To understand the layerwise scaling behaviour of SAM under this scaling, we define the notion of vanishing perturbations .

Vanishing perturbations. The weight perturbation ε l perturbs the l -th layer's activations as

where ˜ δx l denotes the perturbation of the l -th layer's activations accumulated from the weight perturbations { ε l ′ } l ′ ∈ [ l ] in all previous layers. We say a layer l has vanishing perturbations if ˜ δx l → 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 ∥ ε l ′ ∥ ∗ / ∥ W l ′ ∥ ∗ → 0 for all l ′ ∈ [ l ] .

64

256

1024

4096

16384

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

1

.

4

‖

˜

W

-

W

‖

F

/

‖

W

‖

∗

×

10

-

5

Input layer

SAM-SGD

SGD

64

256

1024

4096

16384

0

.

0000

0

.

0002

0

.

0004

0

.

0006

0

.

0008

0

.

0010

0

.

0012

0

.

0014

Hidden layer

64

256

1024

4096

16384

0

.

000

0

.

005

0

.

010

0

.

015

0

.

020

0

.

025

Output layer

64

256

1024

4096

16384

0

.

12

0

.

14

0

.

16

0

.

18

0

.

20

‖

x

t

-

x

0

‖

2

, normalized

64

256

1024

4096

16384

0

.

15

0

.

20

0

.

25

0

.

30

64

256

1024

4096

16384

3

4

5

6

7

8

SAM-SGD

SGD

Informally, Proposition 2 below shows that for every choice of a decay parameter d > 0 , either the training dynamics of SAM are unstable or all the hidden layers of the network have vanishing perturbations in the limit. The formal results are stated in the next section.

Proposition 2 ( Global perturbation scaling is unstable or induces vanishing perturbations ) . Fix ρ > 0 and t ∈ N . Let ˚ f t 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 ρn -d for t steps. If d < 1 / 2 , then output perturbations blow up, and ˚ f t is unstable. If d > 1 / 2 , then the perturbations in all layers vanish and ˚ f t corresponds to the limit after t steps of SGD. If d = 1 / 2 , then only the last layer is effectively perturbed, all other layers have vanishing perturbations.

Figure 2 shows statistics of an MLP trained with (SAM) with global width-dependent scaling ρn -1 / 2 versus the same MLP trained with SAM where only the last-layer weights are perturbed and ε l = 0 for all l ∈ [ 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 3 ( 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 implies 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, it may be undesirable for optimal generalization if only the last layer is perturbed (Figure 1). ◀

Layerwise perturbation scaling. In the next section, we show that correcting the (SAM) update rule to achieve effective perturbations in every single layer requires introducing additional hyperparameters - layerwise width-dependent scaling of the perturbation radius. This is similar in spirit to µ P which corrects standard parameterization by introducing layerwise scaling of the learning rates. We postulate that achieving width-independent scaling of both updates and perturbations is a 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 ( µ P 2 ) that achieves both feature learning and effective perturbations in all layers in the infinite-width limit .

Figure 1 shows that µ P 2 indeed achieves hyperparameter transfer in the optimal joint choice of ( η, ρ ) , while also achieving the best generalization performance (Table 2).

General perturbation scaling condition. For intuitive understanding and a generalization to other perturbation rules, a simple condition for achieving effective perturbations in any layer follows from our results: in every layer, perturbations should scale like updates in µ P .

The reason is that both updates and perturbations are gradient-based ∇ W l L = ∇ h l L· ( x l -1 ) ⊤ , and thus low-rank and correlated with the incoming activations x l -1 . 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 l t ∈ R fan_out × fan_in , its update δW l t and its perturbation ε l t , it should hold at all times t that

with big-O notation that only tracks dependence on network width (Definition C.1). We discuss the spectral perspective in more detail in 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 a new class of parameterizations, which we refer to as bcd -parameterizations . Motivated by the analysis in Section 3, the class of bcd -parameterizations naturally extends abc -parameterizations (Yang and Hu, 2021) by including layerwise scaling of the perturbation radius. By setting all weight multiplier exponents a l = 0 , we do not need to modify the MLP architecture and recover representatives of each abc -parameterization that capture their essence and condense all equations: Ignoring numerical considerations (Blake et al., 2024), each abc -parameterization is essentially a layerwise initialization and learning rate scaling. The effects of weight multipliers on 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 bcd -parameterization, we utilize the theoretical framework of NE ⊗ OR ⊤ programs (Yang et al., 2023b). We write the two forward and backward passes for each SAM update (ascent/perturbation step then descent/update step) using the NE ⊗ OR ⊤ computation rules and rigorously track all relevant scalings as provided by the NE ⊗ OR ⊤ 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, we present our 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 or σ -gelu for σ > 0 sufficiently small. For small enough σ > 0 , σ -gelu (Definition C.9) approximates ReLU arbitrarily well. We also assume constant training time as width n →∞ . 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/perturbation and the descent/update step, as 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 4 ( bcd -parametrization ) . A bcd -parametrization { b l } l ∈ [ L +1] ∪ { c l } l ∈ [ L +1] ∪ { d l } l ∈ [ L +1] ∪ { d } defines the training of an MLP with SAM in the following way:

(b) Train the weights using the SAM update rule with layerwise learning rates,

(a) Initialize weights iid as W l ij ∼ N (0 , n -2 b l ) .

with the scaled perturbation ε t via layerwise perturbation radii,

W.l.o.g. we set ∥ v t ∥ = Θ(1) , which prevents nontrivial width-dependence from the denominator. This imposes the constraints: d 1 ≥ 1 / 2 -min( b L +1 , c L +1 ) , d l ≥ 1 -min( b L +1 , c L +1 ) for l ∈ [2 , L ] , and d L +1 ≥ 1 / 2 , with at least one equality required to hold (see Appendix E.1.3). The normalization v t / ∥ v t ∥ removes one degree of freedom from { d l } l ∈ [ L +1] via the equivalence { d ′ l } l ∈ [ L +1] ∼ = { d l } l ∈ [ L +1] iff there exists a C ∈ R such that d ′ l = d l + C for all l ∈ [ L +1] . ◀

Stability. To ensure that the training dynamics of SAM are well-behaved with scale, we require bcd -parameterizations to satisfy conditions of stability. Perturbed weights ˜ W l = W l + ε l induce perturbed activations x l + ˜ δx l and a perturbed output function ˜ f t ( ξ ) := f ˜ W t ( ξ ) . We call a bcd -parameterization stable (Definition C.3) if the hidden activations have width-independent scaling Θ(1) at initialization and during training, and neither the updates nor the perturbations ˜ δx l of the activations or output logits ˜ f t -f t 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 ˜ r of a bcd -parameterization as

̸

Similar to the maximal feature update scaling r from Yang and Hu (2021), ˜ r describes how much the last hidden-layer activations are perturbed as a function of width, x L + ˜ δx L = Θ( n -˜ r ) . Hidden-layer activation perturbations do not explode with width if and only if ˜ r ≥ 0 . The output perturbations not to blow up if and only if d + d L +1 ≥ 1 and b L +1 + ˜ r ≥ 1 . In particular, this implies that any stable bc -parameterization together with naive perturbation scaling d l = d = 0 for all l ∈ [ 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 width-independent perturbation scaling which is attained iff ˜ r = 0 .

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 a stable bcd -parametrization is perturbation nontrivial if there exists a training routine, t ∈ N 0 and ξ ∈ R d in such that ˜ δf t ( ξ ) := f ˜ W t ( ξ ) -f W t ( ξ ) = Ω(1) . Otherwise, the bcd -parametrization is perturbation trivial . ◀

Theorem 6 ( Perturbation nontriviality characterization ) . A stable bcd -parametrization is perturbation nontrivial if and only if d + d L +1 = 1 or min( b L +1 , c L +1 ) + ˜ r = 1 .

For the class of stable and perturbation nontrivial bcd -parameterizations, 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-width limit? To address this question, we make the following crucial distinction between non-vanishing and effective perturbations .

Non-vanishing versus effective perturbations. Recall that the weight perturbation ε l perturbs the l -th layer's activations as

where ˜ δx l denotes the perturbation of the l -th layer's activations accumulated from the weight perturbations { ε l ′ } l ′ ∈ [ l ] in all previous layers. Therefore, perturbations ˜ δx l can stem both from weight perturbations ε l ′ in a previous layer l ′ < l and/or from weight perturbations ε l in the current layer l . Intuitively, if we perturb a 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 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 l -th layer perturbations are achieved if and only if weight perturbations scale like the weights in spectral norm, ∥ ε l ∥ ∗ / ∥ W l ∥ ∗ = Θ(1) . Without an effective perturbation ε l of the l -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 is crucial.

Definition 7 ( Non-vanishing perturbations ) . For l ∈ [ L ] , we say that a stable parameterization has non-vanishing perturbations in the l -th layer if there exists a t ∈ N such that ˜ δx l t = Ω(1) . ◀

Definition 8 ( Effective perturbations ) . For l ∈ [ L + 1] , we say that a stable parameterization effectively perturbs the l -th layer if there exists a t ∈ N such that ε l t ( x l -1 t + ˜ δx l -1 t ) = Θ(1) , where x 0 t + ˜ δx 0 t = ξ t . ◀

In SP/NTP:

In μP:

I+Ip+p

+Ip+1

十

d

Collapse to SGD

Collapse to SGD

1.0

1.0

up2

Non-trivial

Non-trivial

SAM dynamics

SAM dynamics

Unstable

Unstable

-0.5

0:5

-0.5

0.5

= min(bL+1,CL+1) + d + min(d - I(l ≠ 1))

L

=min(bL+1,CL+1) + d + min(d - I(l ≠ 1))

L

l=1

1=1

Theorem 9 provides a characterization of stable bcd -parameterizations with vanishing perturbations in any given layer.

Theorem 9 ( Vanishing perturbation characterization ) . For any l 0 ∈ [ L ] , the following statements are equivalent:

(a) A stable bcd -parametrization has vanishing perturbations in layer l 0 .

(c) ˜ r l 0 := min( b L +1 , c L +1 ) + d +min l 0 m =1 ( d m -I ( m = 1)) > 0 .

(b) A stable bcd -parametrization has vanishing perturbations in layer l for all 1 ≤ l ≤ l 0 .

̸

It follows from Theorem 9 that any stable bcd -parameterization that performs updates in the original gradient direction (i.e., d l = C for all l ∈ [ L +1] for some C ∈ R ) has vanishing perturbations in all input and hidden layers l ∈ [ L ] , and the last layer l = L +1 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 9 implies that a stable bcd -parameterization has non-vanishing perturbations in layer l 0 if and only if ˜ r l 0 = 0 . Achieving effective perturbations is a stronger requirement for which Theorem 10 provides the necessary and sufficient conditions.

Theorem 10 ( Effective perturbation characterization ) . For l ∈ [ L ] , a stable bcd -parametrization effectively perturbs the l -th layer if and only if min( b L +1 , c L +1 ) + d + d l -I ( l = 1) = 0 .

A stable bcd -parametrization effectively perturbs the last layer if and only if d + d L +1 = 1 .

4.2 Maximal Update and Perturbation Parameterization ( µ P 2 )

We postulate that just as the optimal learning rate transfers across widths under µ P for SGD and Adam due to non-vanishing width-independent feature evolution in all layers, the optimal learning rate and perturbation radius may be jointly transferable across widths if additionally the weight perturbations induce width-independent perturbations of the activations in all layers. Here, we show that, for every stable initialization and learning rate scaling with b L +1 ≥ 1 , there exists a unique stable layerwise perturbation scaling which effectively perturbs every single layer. We term this layerwise perturbation scaling { d l } l ∈ [ L +1] ∪ { 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 bcd -parametrization { b l } l ∈ [ L +1] ∪ { c l } l ∈ [ L +1] ∪ { d l } l ∈ [ L +1] ∪ { d } . If b L +1 < 1 , then there does not exist a stable choice of { d l } l ∈ [ L +1] ∪ { d } that achieves effective perturbations before the last layer. If b L +1 ≥ 1 , then up to the equivalence d ′ l = d l + C , C ∈ R , ∀ l ∈ [ L +1] , the unique stable choice { d l } l ∈ [ L +1] ∪ { d } that effectively perturbs all layers l ∈ [ L +1] is given by

̸

Maximal Update and Perturbation Parameterization µ P 2 . To achieve feature learning in every layer and hyperparameter transfer in the learning rate, µ P is the unique 2 choice of layerwise initialization variance and learning rate scalings { b l , c l } l ∈ [ L +1] (Yang and Hu, 2021). Together with Theorem 11, this shows that there exists a unique 2 bcd -parameterization that achieves both feature learning and effective perturbations in all layers, we call maximal update and perturbation parametrization , µ P 2 for short. Now that we have found a parameterization that achieves widthindependent scaling of both activation updates and activation perturbations, µ P 2 fulfills essential necessary conditions for hyperparameter transfer to occur in both η and ρ .

Remark 12 ( Achieving µ P 2 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 { a l } l ∈ [ L +1] , there is a unique 2 choice of bcd -hyperparameters 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 ∥∇ W L∥ F in (SAM). The term ∥∇ W L∥ F 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 -a l · W l with

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 -µP 2 ), that also applies to naive weight multipliers a l = 0 , is its incompatibility with unit scaling considerations for low precision training (Blake et al., 2024). ◀

Alternative perturbation scaling definitions. Scaling equivalent to ( a -µP 2 ) 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-independent (see perturbation rule (DP) and Appendix F.7 for more details). The ablations in Appendix H.4 suggest that this has a negligible effect on the optimal generalization performance of µ P 2 , 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 b L +1 ≥ 1 is required for stable feature learning. Theorem 11 shows that b L +1 ≥ 1 is also required for effective hidden-layer perturbations. Beyond this condition, the choice of { b l } and { c l } is decoupled from that of perturbation scalings { d l } ∪ { d } for stable bcd -parameterizations, because the scale of the activations of a layer l is entirely determined by the scale of initialization b l and learning rates c l , given stability. Consequently, whether a parameterization is trivial , in the lazy regime, or in the feature learning regime is independent of the choice of d l 's provided that all stability constraints are met. A complete characterization of these regimes for the class of bc -parameterizations has been provided in Yang and Hu (2021) and remains unchanged for the class of stable bcd -parameterizations. For completeness, formal definitions and the corresponding results are stated in Appendices C and D.

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, including all ResNet and Transformer components (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 .

Generalization to other SAM variants. We would like to find the correct layerwise perturbation scaling without writing out the NE ⊗ OR ⊤ program 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 µ P if and only if ε l = Θ( δW l ) . This condition holds as soon as weight updates δW l and perturbations ε l are both correlated with the incoming activations x l -1 , for example if both are gradient-based. Table 1 summarizes the application of this condition to two ASAM variants that perform well empirically but cannot be written as a NE ⊗ OR ⊤ 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

µ P 2 , it holds that ε l = Θ( n -1 ) but W l = Θ( n -1 / 2 ) entrywise, due to large initialization, showing that it is crucial to compare perturbations to updates or to measure weight scalings in spectral norm.

5 The maximal update and perturbation parameterization µ P 2 achieves hyperparameter transfer and improved generalization

In this section, we provide experimental results showing that µ P 2 achieves hyperparameter transfer in both η and ρ across architectures, and that µ P 2 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) on CIFAR10 (Krizhevsky et al., 2009) and Vision Transformers (ViTs) (Dosovitskiy et al., 2021) on Imagenet1K (Deng et al., 2009). While we directly implement bcd -parameterizations for MLPs and ResNets in PyTorch (Paszke et al., 2019), we use the mup -package (Yang et al., 2022) as a basis for ViT experiments. Pseudocode and a spectral derivation of our µ P 2 -implementation for ViTs, which is equivalent to ( a -µP 2 ), 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 µ P-literature mostly focuses on the more immediate transfer in training error, for SAM it is crucial to consider optimality in test error as the perturbation radius acts as a regularizer, so that optimality in test error typically coincides with suboptimal training error. In µ P without perturbation scaling, the regime of stable perturbation radii shrinks (Figure H.6). In µ P with global perturbation scaling ρ · n -1 / 2 , the regime of stable ρ remains invariant under width scaling, but there is no significant improvement of SAM beyond SGD, so that the optimal perturbation radius fluctuates within its stable regime due to noise. Only µ P 2 consistently achieves hyperparameter transfer across widths, and achieves significant improvement over its base optimizer SGD in µ P at scale. The full hyperparameter landscapes are provided in Appendix H.3.

ρ -transfer in ViTs. Figure 4 shows that the optimal perturbation radius transfers for ViT-S/16 on Imagenet1K trained with SAM in µ P 2 . While Andriushchenko and Flammarion (2022, Appendix E.3) observe diminishing benefits of SAM at large widths in SP, here the improvements beyond the base optimizer AdamW in µ P are particularly large.

ρ -transfer for SAM variants in µ P 2 . Figure 6 shows that training a ResNet-18 in µ P 2 achieves hyperparameter transfer in ρ for all considered SAM variants with varying width. µ P with global perturbation scaling ( µ Pglobal) has a width-invariant stability threshold in ρ and the optimal ρ clearly shifts toward that threshold. It would be interesting to see whether this shift continues with larger width and leads to suboptimal performance of µ P-global in wider ResNets. Table 2 shows that all SAM variants perform similarly well in µ P 2 ,

10

-

2

10

-

1

10

0

10

1

Perturbation radius

ρ

0

.

48

0

.

50

0

.

52

0

.

54

0

.

56

0

.

58

0

.

60

0

.

62

Top 1 accuracy

width

768

384

192

some slightly outperforming the best-

10

-

2

10

-

1

10

0

Perturbation radius

ρ

94

95

96

97

Test accuracy

SAM

µ

P-global

width

4.0

2.0

1.0

0.5

10

-

2

10

-

1

10

0

Perturbation radius

ρ

94

95

96

97

SAM

µ

P

2

10

-

1

10

0

10

1

Perturbation radius

ρ

94

95

96

97

SAM-ON

µ

P

2

10

-

1

10

0

10

1

Perturbation radius

ρ

94

95

96

97

Elem. ASAM

µ

P

2

performing variant SAM-ON in SP. This suggests that for ResNets, even with a proper layerwise 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. Table 1 ( left ) explains these observations by showing that only these two SAM variants effectively perturb normalization layers under global perturbation scaling. Table 1 ( right ) also provides full control over which layers to perturb. For transferring the optimal ρ with SAM-ON in µ P, our theory predicts the global scaling ρ = Θ( n 1 / 2 ) which is confirmed by our empirical observations (Figure 6). However, properly understanding the role of normalization layer perturbations remains an important question for future work. Note that we report results after fine-tuning all hyperparameters. The performance gain of µ P 2 over

0

25

50

75

100

125

150

175

200

Epochs of Training

10

1

10

2

100% - test accuracy

SAM

µ

P

2

SGD

µ

P

SAM SP

SGD SP

SP and µ P-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, µ P 2 appears to stabilize SAM's training dynamics compared to SP (Figure 5).

6 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, we observe that MLPs and ResNets in SP can sometimes display HP transfer in η and ρ after multi-epoch training to convergence (Appendix H.3.2). This goes beyond the observations by Everett et al. (2024) as we observe transfer even without tuning layerwise learning 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.

References

Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. In International Conference on Machine Learning (ICML) , 2022. Cited on 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:2305.16292 , 2023a. Cited on page 38, 53.

Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flammarion. 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-aware minimization: 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 Barnett, 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µ P: The unit-scaled maximal update parametrization. arXiv: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 on 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) , 33:1877-1901, 2020. Cited on page 60.

Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. arXiv:2106.01548 , 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. In 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:2401.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. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2009. Cited on page 9, 51, 71.

Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Proceedings of the 34th International Conference on Machine Learning (ICML) , 2017. Cited on page 17.

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 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR) , 2021. Cited on page 9.

Katie E Everett, Lechao Xiao, Mitchell Wortsman, Alexander A Alemi, Roman Novak, Peter J Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie 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 International 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 on 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 9.

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-optimal large language models. arXiv:2203.15556 , 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.12226 , 2023. Cited on page 17.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and generalization 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. Fantastic generalization measures and where to find them. In International Conference on Learning Representations (ICLR) , 2020. Cited on page 17.

Jean Kaddour, Linqing 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, Alec 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, et 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 on page 1, 3, 17, 37, 39, 48.

Yann LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robert 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 log-gaussian: Resnets and their infinite-depth-andwidth 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/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , 2022. Cited on page 3, 17.

Philip M. Long and Peter L. Bartlett. Sharpness-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 twolayer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665-E7671, 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 1, 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 in 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-width 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 DeVito, 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 in deep neural networks through transient chaos. Advances in Neural Information Processing Systems (NeurIPS) , 29, 2016. Cited on page 16.

David Samuel. (Adaptive) SAM Optimizer (PyTorch). https://github.com/davda54/sam , 2022. Cited on page 37, 48, 49, 71.

Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. arXiv:1611.01232 , 2016. Cited on page 16.

Sungbin Shin, Dongyeop Lee, Maksym Andriushchenko, and Namhoon Lee. The effects of overparameterization 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 Thirty-eighth Annual Conference on Neural Information Processing Systems (NeurIPS) , 2024. Cited on page 17.

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 generalization in deep neural networks. In International Conference on 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 2, 3, 16, 42.

Greg Yang. Tensor programs iii: 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, 3, 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. arXiv:2203.03466 , 2022. Cited on page 2, 3, 9, 16, 19, 43, 51, 57, 60, 61, 71.

Greg Yang, James B. Simon, and Jeremy Bernstein. A 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, 5, 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.

Appendices

Appendix Contents.

A Notation

Table A.1: (Notation) Overview over notation used in the main paper (top) and in the appendix (bottom).

B Detailed related work

Signal propagation. Our work can be seen as scaling theory with the goal of preventing both vanishing and exploding signals in forward and backward passes, where the analysis of SAM requires considering stability of perturbations in each layer as well. In this sense, we build on a rich literature, often restricted to an analysis at initialization (Schoenholz et al., 2016; Poole et al., 2016; Hanin and Rolnick, 2018; Xiao et 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 1 √ L -scaling 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.

Tensor Programs. After kernel-based approaches to understand infinite-width limits of neural networks (Neal, 1996; Jacot et al., 2018) and applications of mean-field theory (Mei et 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 abc -parameterizations, where each abc -parameterization 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 ( µ P) 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 a global learning rate η > 0 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 µ P (Wenger et al., 2023; Vyas et al., 2024). Yang and Littwin (2023) characterize stable learning with adaptive optimizers at infinite width into a feature learning versus a (nonlinear) operator regime. SAM

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 µ P is equivalent to the spectral scaling conditions on the weights ∥ ∆ W l ∥ = Θ( √ n l /n l -1 ) and ∥ ∆ W l ∥ = Θ( √ n l /n l -1 ) . 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 √ fan_out / fan_in 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 et 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-norm adaptive edge of stability at which it 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 a 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 l w ( x ) = | W l | ⊙ x and layerwise normalization T l w ( x ) = ∥ W l ∥ F · x . 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 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. 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 architecture-agnostic analysis, they show that 1-SAM minimizes the trace of Hessian like average-case sharpness, for small enough η and ρ . Similarly, the theoretical results by Andriushchenko and Flammarion (2022) rely on the assumption that learning rate η and perturbation radius ρ are chosen sufficiently close to 0 . Arguably, the empirically optimal choice of η and ρ 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-width limit is not only of theoretical interest but entails immediate practical benefits.

Liu et al. (2022) introduce Look-LayerSAM 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

these kinds of algorithms follows

with some ϕ : R + → R + and where ∇ W l L may be replaced by Adam's m t √ v t + ε . In practice, often simple functions like ϕ ( x ) = 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. A derivation of µ P for LAMB could also yield feature learning in all layers in the infinite-width limit as well as hyperparameter transfer. It certainly requires layerwise learning rate scaling. In the case ϕ ( x ) = x , following a heuristic scaling derivation as in Appendix F.4 leads to layerwise learning rate scalings η 1 = η L +1 = Θ(1) and η l = Θ( n -1 / 2 ) for hidden layers l ∈ [2 , L ] . With a bounded function like ϕ ( x ) = max( c, min( x, C )) , the scalings become η 1 = Θ( n 1 / 2 ) , η L +1 = Θ( n -1 / 2 ) and η l = Θ(1) for hidden layers l ∈ [2 , L ] . We leave a closer investigation of feature learning and hyperparameter transfer with LAMB and Look-LayerSAM in SP and µ P to future work.

C Definitions

In this section, we collect all definitions that do not 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 to width n →∞ .

scalar random variable sequence {√ ∥ x n ∥ 2 /n } n is Θ( n -a ) . For the definition of c = O ( n -a ) -a -a -a -a

Definition C.1 ( Big-O Notation ) . Given a sequence of scalar random variables c = { c n ∈ R } ∞ n =1 , we write c = Θ( n -a ) if there exist constants A,B such that for almost every instantiation of c = { c n ∈ R } ∞ n =1 , for n large enough, An -a ≤ | c n | ≤ Bn -a . Given a sequence of random vectors x = { x n ∈ R n } ∞ n =1 , we say x has coordinates of size Θ( n -a ) and write x = Θ( n -a ) to mean the and

c

with

= Ω(

c

n

n

Bn almost surely.

|

, adapt the above definition of

)

-

a

and

An

-

a

c

≤ |

|

·

∥

∥

Bn

/n

→

0

◀

Definition C.2 ( Training routine ) . A training routine is a combination of base learning rate η ≥ 0 , perturbation radius ρ ≥ 0 , training sequence { ( ξ t , y t ) } t ∈ N and a continuously differentiable loss function L ( f ( ξ ) , 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 a bcd -parametrization of an L -hidden layer MLP is stable if

For every nonzero input ξ ∈ R d in \ { 0 } ,

where the expectation is taken over the random initialization.

For any training routine, any time t ∈ N , l ∈ [ L ] , ξ ∈ R d in , we have

where the hidden constant in O ∗ can depend on the training routine, t , ξ , l and the initial function f 0 .

For any training routine, any time t ∈ N 0 , l ∈ [ L ] , ξ ∈ R d in , for the perturbed (pre-)activation ˜ h l t := h l ( ˜ W t ) , ˜ x l t := x l ( ˜ W t ) and output function ˜ f t ( ˜ W t ) we have

where the hidden constant in O ∗ can depend on the training routine, t , ξ , l and the initial function f 0 .

◀

n

c

= Θ(

n

)

, respectively. We write by replacing

x

n

=

o

(

n

-

a

)

An if

n

≤ |

c

n

x

n

a

√

2

Definition C.4 ( Nontriviality ) . We say a bcd -parametrization is trivial if for every training routine, f t ( ξ ) -f 0 ( ξ ) → 0 almost surely for n →∞ , for every time t > 0 and input ξ ∈ R d in . Otherwise the bcd -parametrization is nontrivial . ◀

Definition C.5 ( Feature learning ) . We say a bcd -parametrization admits feature learning in the l -th layer if there exists a training routine, a time t > 0 and input ξ such that x l t ( ξ ) -x l 0 ( ξ ) = Ω ∗ (1) , where the constant may depend on the training routine, the time t , the input ξ and the initial function f 0 but not on the width n . ◀

Definition C.6 ( Vanishing perturbations ) . Let l ∈ [ L ] . We say that a stable bcd -parametrization has vanishing perturbations in the l -th layer if for any training routine, t ∈ N 0 and ξ ∈ R d in , it holds that ˜ x l t -x l t = o (1) , and it has vanishing perturbations in the output if for any training routine, t ∈ N 0 and ξ ∈ R d in it holds that ˜ δf t ( ξ ) := f ˜ W t ( ξ ) -f W t ( ξ ) = o (1) . ◀

Definition C.7 ( Perturbation nontriviality ) . Let l ∈ [ L ] . We say that a stable bcd -parametrization is perturbation nontrivial with respect to the l -th layer if and only if it does not have vanishing perturbations in the l -th layer. A stable bcd -parametrization is perturbation nontrivial with respect to the output if it does not have vanishing perturbations in the output. ◀

Definition C.8 ( Effective perturbations ) . Let l ∈ [ L +1] . We say that a stable bcd -parametrization effectively perturbs the l -th layer if there exists a training routine, t ∈ N and ξ ∈ R d in such that ˜ δW l t ˜ x l -1 t ( ξ ) = Θ(1) where ˜ δW l t is defined in (LP) and ˜ x 0 t = x 0 t = ξ t . ◀

In order to apply the Tensor Program Master Theorem, all Nonlin and Moment operations in the NE ⊗ OR ⊤ program, which do not only contain parameters as inputs, are required to be pseudoLipschitz in all of their arguments. For training with SGD, this is fulfilled as soon as ϕ ′ is pseudoLipschitz. Both tanh as well as σ -gelu fulfill this assumption.

Definition C.10 ( Pseudo-Lipschitz ) . A function f : R k → R is called pseudo-Lipschitz of degree d if there exists a C > 0 such that | f ( x ) -f ( y ) | ≤ C ∥ x -y ∥ (1 + ∑ k i =1 | x i | d + | y i | d ) . We say f is pseudo-Lipschitz if it is so for any degree d . ◀

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 NE ⊗ OR ⊤ program, scalings and infinite-width limits.

Recall that our definition of bcd -parameterizations extends abc -parameterizations by setting the maximal perturbation scaling to n -d and allowing relative downweighting n -d l of the global scaling in each layer l . The perturbation scaling does not affect the choice of layerwise initialization variance scalings b l and the layerwise learning rate scalings c l . Common bc -parametrizations for SGD are summarized in Table D.1. SAM with SGD as a 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, we require width-independent denominator scaling ∥ v t ∥ = Θ(1) of the scaled gradient for the perturbation (LP), which imposes the constraints

All (pre-)activation and function outputs can be thought of as outputs given a fixed input ξ ∈ R d in \ { 0 } with d in ∈ N fixed, e.g. f t := f W t := f W t ( ξ ) . For the perturbed weights we write ˜ W t := W t + ˜ δW t , with ˜ δW t defined in (LP) as ε l t . Here we write weight perturbations as ˜ δW l t instead of ε l t to show the resemblance to weight updates δW l t . Perturbed activations and function outputs at time t are written as ˜ x l t ( ξ ) = x l ˜ W t ( ξ ) and ˜ f t ( ξ ) = f ˜ W l t ( ξ ) . 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 is either tanh or σ -gelu for σ > 0 sufficiently small. ◀

We define the maximal feature update scale of a bcd -parameterization

̸

as well as the maximal feature perturbation scale of a bcd -parameterization

̸

Stability requires the constraints (a-c) from SGD and additional perturbation stability constraints (d-e) that include the layerwise perturbation scales { d l } l =1 ,...,L +1 .

Theorem D.2 ( Stability characterization ) . A bcd -parametrization is stable if and only if all of the following are true:

(a) (Stability at initialization, h l 0 , x l 0 = Θ(1) for all l , f 0 = O (1) ) b 1 = 0 , b l = 1 / 2 for l ∈ [2 , L ] and b L +1 ≥ 1 / 2 .

(c) (Output function does not blow up during training, i.e. ∆ W L +1 t x L t , W L +1 0 ∆ x L t = O (1) ) c L +1 ≥ 1 and b L +1 + r ≥ 1 .

(b) (Features do not blow up during training, i.e. ∆ x l t = O (1) for all l ) r ≥ 0 .

(d) (Feature perturbations do not blow up, i.e. δx l t = O (1) for all l )

(e) (Output function perturbations do not blow up during training, i.e. ˜ δW L +1 t ˜ x L t , W L +1 t ˜ δx L t = O (1) )

˜ ˜ r ≥ 0 .

d

+

d

L

+1

1

and

b

L

+1

+ ˜

r

1

.

The nontriviality and feature learning characterizations from SGD remain unaltered. This is because in the definition of r , it holds that d + d L +1 ≥ 1 (from perturbation stability), and min( b L +1 , c L +1 ) ≤ 1 already had to hold for nontriviality in SGD, so that stable perturbation scaling does not affect r .

Theorem D.3 ( Nontriviality characterization ) . A stable bcd -parametrization is nontrivial if and only if c L +1 = 1 or min( b L +1 , c L +1 ) + r = 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 { d l } l ∈ [ L +1] ∪ { d } . Hence the conditions are the same as for SGD. Below we provide a slightly refined result in terms of the maximal feature update scale r l 0 of a bcd -parameterization up to layer l 0 (as provided in the Appendix of Yang and Hu (2021)).

Theorem D.4 ( Feature learning characterization ) . For any l 0 ∈ [ L ] , the following statements are equivalent:

(a) A stable, nontrivial bcd -parametrization admits feature learning in layer l 0 .

(c) r l 0 := min( b L +1 , c L +1 , d + d L +1 ) + min l 0 m =1 ( c m -I ( m = 1)) = 0 .

̸

(b) A stable, nontrivial bcd -parametrization admits feature learning in layer l for all l ≥ l 0 .

Consequently, a stable, nontrivial bcd -parametrization admits feature learning (at least in the last layer activations) if and only if r = 0 .

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 l -th layer as δW l t x l -1 t = Θ(1) and it would occur if and only if min( b L +1 , c L +1 , d + d L +1 ) + c l -I ( l = 1) = 0 . ◀

̸

As for nontriviality, perturbation nontriviality in the output is attained if the constraints for ˜ δW L +1 t ˜ x L t or W l t ˜ δx L t are exactly satisfied.

Theorem D.6 ( Perturbation nontriviality characterization ) . Let l ∈ [ L ] . A stable bcd -parametrization is perturbation nontrivial with respect to the l -th layer if and only if

̸

A stable bcd -parametrization is perturbation nontrivial with respect to the output if and only if d + d L +1 = 1 or min( b L +1 , c L +1 ) + ˜ r = 1 .

The converse formulation of the perturbation-nontriviality results characterizes the regime of vanishing perturbations.

Corollary D.7 ( Vanishing perturbation characterization ) . For any l 0 ∈ [ L ] , the following statements are equivalent:

(a) A stable bcd -parametrization has vanishing perturbations in layer l 0 .

(c) ˜ r l 0 := min( b L +1 , c L +1 ) + d +min l 0 m =1 ( d m -I ( m = 1)) > 0 .

̸

(b) A stable bcd -parametrization has vanishing perturbations in layer l for all 1 ≤ l ≤ l 0 .

A stable bcd -parametrization has vanishing perturbations with respect to all layers and the output function if and only if d L +1 > 1 / 2 and ˜ r > max(0 , 1 -b L +1 ) . This case reduces to 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 l ∈ [ L ] , a stable bcd -parametrization effectively performs SAM in the l -th layer if and only if min( b L +1 , c L +1 ) + d + d l -I ( l = 1) = 0 . A stable bcd -parametrization effectively performs SAM in the last layer if and only if d + d L +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 ∥ v t ∥ can be approximated efficiently without an additional backward pass under global perturbation scaling.

Theorem D.9 ( Global Perturbation Scaling ) . Given any stable bcd -parametrization { b l } l ∈ [ L +1] ∪ { c l } l ∈ [ L +1] ∪ { d l } l ∈ [ L +1] ∪ { d } . The parametrization performs updates in the original gradient direction if and only if d l = C for all l ∈ [ L +1] for some C ∈ R . In this case, the parametrization has vanishing perturbations in all hidden layers l ∈ [ L ] , and the last layer l = L +1 is effectively perturbed if and only if d = 1 / 2 . If b L +1 > 1 / 2 (as in µ P), the gradient norm is dominated by the last layer and simplifies to,

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 bcd -parameterizations. Of course, if we add even more hyperparameters by decoupling numerator and denominator scalings, we can set all contributions to Θ(1) , which is what we do in Appendix F.7.

Proposition D.10 ( Balancing gradient norm contributions ) . Given any stable bcd -parametrization { b l } l ∈ [ L +1] ∪ { c l } l ∈ [ L +1] ∪ { d l } l ∈ [ L +1] ∪ { d } . If all layers contribute to the gradient norm nonvanishingly in the limit, i.e. ∥ v l t ∥ = Θ( ∥ v t ∥ ) for all l ∈ [ L +1] , t ∈ N 0 , then the parametrization has vanishing perturbations in all hidden layers l ∈ [ L ] . Such a parametrization effectively performs SAM in the last layer l = L +1 if and only if d = 1 / 2 .

The following theorem provides the unique correct perturbation scaling for any stable bc -parameterization with b L +1 ≥ 1 .

Theorem D.11 ( Perturbation Scaling Choice for Effective Perturbations ) . Given any stable bcd -parametrization { b l } l ∈ [ L +1] ∪ { c l } l ∈ [ L +1] ∪ { d l } l ∈ [ L +1] ∪ { d } . If b L +1 < 1 , then there does not exist a stable choice of { d l } l ∈ [ L +1] ∪ { d } that achieves effective perturbations before the last layer. If b L +1 ≥ 1 , then up to the equivalence d ′ l = d l + C , C ∈ R , ∀ l ∈ [ L +1] , the unique stable choice { d l } l ∈ [ L +1] ∪ { d } with effective perturbations in all layers l ∈ [ L +1] is given by

In this parameterization, the first layer dominates the gradient norm as

Table D.2: (Perturbation scalings) Overview over important choices of the global perturbation scaling ρn -d and the layerwise perturbation scalings n -d l 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 c ∇ := min( b L +1 , c L +1 ) .

Table D.2 summarizes the consequences of Theorem D.11. Together with Theorem D.11, the following proposition suggests that b L +1 = 1 is a good choice. However b L +1 > 1 can also induce effective perturbations, as long as d and d L +1 are chosen correctly.

Proposition D.12 ( Effects of last-layer initialization b L +1 on all perturbations ) . If a stable bcd -parametrization with min( b L +1 , c L +1 ) ≤ 1 is perturbation nontrivial with respect to any hidden layer l ∈ [ 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 bcd -parametrizations with d 1 = -min( b L +1 , c L +1 ) -d 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 bcd -parametrization, we can compute the maximal layer l 0 such that ˜ r l 0 > 0 , and in wide networks do not have to compute SAM perturbations before layer l 0 + 1 ; as soon as b L +1 > 1 / 2 (as for µ P), the gradient norm for the SAM update rule is approximately given by ∥∇ L t ∥ ≈ L ′ ( f t ( ξ t ) , y t ) ∥ x L t ∥ , which can directly be computed without an additional backward pass. The practical recommendation from our experiments however is to either use µ P 2 or 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

adapted NE ⊗ OR ⊤ program and bcd -constraints. Note that standard parametrization gets even more unstable without dividing by ∥∇ L ∥ = Θ( n 1 / 2 ) , now requiring d L +1 ≥ 1 for stability. Similar to the previous results, this shows that unawareness of bcd -parametrizations requires strongly scaling down ρ 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 NE ⊗ OR ⊤ program that corresponds to training a MLP without biases with SAM. For simplicity and clarity of the proof, we prove the one-dimensional case d in = 1 , d out = 1 , but an extension to arbitrary but fixed d in , d out is straightforward. Recall Assumption 1 that allows us to apply the Tensor Program Master Theorem and explicitly state the infinite-width 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 2 0 , . . . , W L 0 as ( W l 0 ) αβ ∼ N (0 , 1 /n ) , which absorbs b l = 1 / 2 .

We initialize the input layer matrix W 1 0 ∈ R n × 1 and normalized output layer matrix ˆ W L +1 0 = W L +1 0 n b L +1 ∈ R 1 × n as ( W 1 0 ) α , ( ˆ W L +1 0 ) α ∼ N (0 , 1) , as initial vectors should have a distribution that is Θ(1) .

In the NE ⊗ OR ⊤ formulation, we write all quantities as θ z z , where θ z denotes their scaling n C for some C ∈ R and z therefore has a Θ(1) distribution. The stability, nontriviality and feature learning conditions then stem from requiring either θ z → 0 or θ z = 1 depending on z and its desired scale.

E.1.2 First forward pass

We denote a definition of a Tensor Program (TP) or NE ⊗ OR ⊤ computation as := . Compared to MLPs trained with SGD nothing changes in the first forward pass,

In the case of MuP, f 0 ( ξ ) = W L +1 0 x L 0 ( ξ ) → 0 defines a scalar in the TP.

Observe the scalings x 1 0 = Θ( h 1 0 ) = Θ( n -b 1 ) , x l 0 = Θ( h l 0 ) = Θ( n 1 / 2 -b l ) for l ∈ [2 , L ] due to CLT, independence at initialization and x l 0 = Θ( h l 0 ) = Θ(1) by stability. Hence stability at initialization inductively requires b 1 = 0 , b l = 1 / 2 for l ∈ [2 , L ] and b L +1 ≥ 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 l t and the normalized perturbations by ˜ δW l t . Before computing the updates we have to compute a full backward pass to determine these perturbed weights for each layer, and then compute a forward pass with these perturbed weights to compute the perturbed preactivations ˜ h l t that we will need for computing the SAM update. Therefore the NE ⊗ OR ⊤ program for SAM maintains a perturbed copy of all preactivations, activations, lastlayer weights and logits just for computing the updates of the actual parameters.

Under MuP, the loss derivative with respect to the function remains χ 0 := L ′ ( f 0 ( ξ 0 ) , y 0 ) → ◦ χ 0 := L ′ (0 , y 0 ) . For the weight perturbation, we need to perform a SGD backward pass, where dz := θ -1 ∇ ∇ z f . For SGD (and for SAM, as we will see later) all gradients have scaling θ ∇ := n -b L +1 in the first step, whereas we overload the notation θ ∇ := n -min( b L +1 ,c L +1 ) for all later steps. For clarity of presentation assume b L +1 ≥ c L +1 here, the other case follows analogously. For the first step this can be understood from

since h L 0 = Θ(1) by the stability assumption, and this scale Θ( n -b L +1 ) propagates through all layers via the chain rule and remains stable in later backward passes. For hidden layer gradients, observe that

where first term's scale stems from the products ( W L 0 ) T W L 0 v = Θ( n 1 -2 b L v ) due to Yang (2021), b L = 1 / 2 for stability at initialization and b L +1 + c L ≥ 1 for update stability during training ( r ≥ 0 ). If 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 l ∈ [2 , L ] is given by the scalar,

where χ t = L ′ ( f t ( ξ t ) , y t ) and we used ∂h l /∂W l ij = ( x l -1 j δ ik ) k =1 ,...,n .

Hence the gradient norm of all weights jointly is given by the unnormalized scalar

with scaling θ 2 ∥∇∥ = Θ( n 2 θ 2 ∇ + n ) = Θ( n ) , because stability at initialization requires b L +1 ≥ 1 / 2 so that n 2 θ 2 ∇ ≤ n . Note that the first layer contributes vanishingly to the gradient norm, the hidden layer gradients only if b L +1 = 1 / 2 (equivalently f 0 = Θ(1) ) and the last-layer activations always in dominating order. So in µ P, in the limit, ∥∇ w L t ∥ = L ′ ( f t ( ξ t ) , y t ) ∥ x L t ∥ . This means that the unscaled gradient always aligns with the last-layer activation. For learning in µ P, this dominance is corrected by the layerwise learning rates.

The squared norm of the rescaled gradient is given by

with scaling θ 2 v = Θ( n 1 -2 d 1 θ 2 ∇ + ∑ L l =2 n 2 -2 d l θ 2 ∇ + n 1 -2 d L +1 ) . For simplicity, set θ v = 1 . This raises the constraints n 1 -2 d 1 θ 2 ∇ ≤ 1 , n 2 -2 d l θ 2 ∇ ≤ 1 for l ∈ [2 , L ] and n 1 -2 d L +1 ≤ 1 , which can be rewritten as

where at least one equality is demanded to hold in order to attain θ v = 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 θ v = 1 the square root of (E.2) defines a normalized TP scalar.

Perturbations. Stability implies that also the perturbed (pre-)activations 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 .

Definition E.1 (Perturbation stability) . We call a bcd -parametrization perturbation stable if and only if ˜ h l t , ˜ x l t = Θ(1) for all l ∈ [ L ] and t ∈ N and ˜ δf t = O (1) for all t ∈ N . ◀

Mathematically we get the normalized weight perturbations for l ∈ { 2 , . . . , L } , which scale as ˜ θ L +1 := ˜ θ W L +1 := n -( d + d L +1 ) , Θ( n ( d + d l ) -b L +1 ) and Θ( n -( d + d 1 ) -b L +1 ) respectively. But the NE ⊗ OR ⊤ program computation rules do not allow to compute matrices ˜ δW l 0 , l ∈ [ L ] , therefore we use the weight updates to directly compute the preactivation and activation changes analogous to the t -th forward pass. For all t ≥ 0 , we write

with the perturbations for l ∈ [2 , L ] , which defines a NonLin operation with the vectors W l 0 ˜ δx l -1 0 and dh l 0 and everything else treated as scalars, and with first backward pass scalings ˜ θ W 1 := n -( d + d 1 ) θ ∇ , ˜ θ W l := n 1 -( d + d l ) θ ∇ and ˜ θ l := max( ˜ θ l -1 , ˜ θ W l ) = max l m =1 ˜ θ W m , where we used that ˜ x l -1 0 = Θ(1) due to perturbation stability. Note that these scalings may implicitly increase when t > 0 since θ ∇ = n -b L +1 gets replaced by θ ∇ = n -min( b L +1 ,c L +1 ) .

The activation perturbations can then simply be defined via the NonLin operation, with the same scaling as ˜ δh l 0 .

The perturbation of the scalar output function can simply be defined via the NonLin operation,

with ˜ θ ′ L +1 := n ˜ θ W L +1 and ˜ θ ′ L ∇ := nθ ∇ ˜ θ L .

SAMUpdate. Finally, we can compute the SAM updates as follows. In the case min( b L +1 , c L +1 ) ≤ d + d L +1 the weight perturbation scale is dominated by the weight scale, so that

with ˜ θ ( L +1) / ∇ := ˜ θ L +1 /θ ∇ ≤ 1 , whereas if min( b L +1 , c L +1 ) > d + d L +1 we write

with θ ∇ / ( L +1) := θ ∇ / ˜ θ L +1 ≤ 1 . In any case, the scaling of dx L SAM, 0 and all other SAM gradients is θ SAM := max( θ ∇ , n -( d + d L +1 ) ) = n -min( b L +1 ,c L +1 ,d + d L +1 ) . The other SAM gradients are given by

where the last line define a NonLin operation in the vectors ( W l 0 ) T dh l SAM,t and x l -1 0 and everything else treated as scalars. Consequently, ∇ h l 0 f | ˜ W 0 is of the same scale as ∇ x l 0 f | ˜ W 0 and ∇ x l -1 0 f | ˜ W 0 is of the scale max( θ SAM , ˜ θ W l θ SAM ) = θ SAM since ˜ θ W l ≤ 1 is required for perturbation stability.

Note that for SAM's weight updates the loss derivative is also evaluated on the perturbed weights,

Constraints on the output function. Assuming ˜ x L 0 = Θ(1) (perturbation stability), we get ˜ χ 0 = O (1) if and only if ˜ δW L +1 0 = O ( n -1 ) if and only if d + d L +1 ≥ 1 .

Hence perturbation stability of the output function holds only if d + d L +1 ≥ 1 and b L +1 + ˜ r L ≥ 1 . Then, perturbation nontriviality holds if and only if d + d L +1 = 1 or b L +1 + ˜ r L = 1 .

We have ˜ χ 0 = Θ(1) if and only if ˜ f 0 = ˜ W L +1 0 ˜ x L 0 = Θ(1) . This can either be caused by changes in the last-layer weights, by non-vanishing initial function W L +1 0 x L 0 (if and only if b L +1 = 1 / 2 ) or by W L +1 0 ˜ δx L 0 = Θ(1) , which holds if and only if b L +1 + ˜ r L = 1 (analogously, W L +1 0 ˜ δx L 0 = O (1) if and only if b L +1 + ˜ r L ≥ 1 ). The first case requires ˜ δW L +1 0 = Θ( n -1 ) , since ˜ δW L +1 0 and ˜ x L 0 are highly correlated. ˜ δW L +1 0 = Θ( n -1 ) is fulfilled if and only if d + d L +1 = 1 (the analogue to c L +1 ≥ 1 for stability and c L +1 = 1 for nontriviality).

In the t -th backward pass, b L +1 + ˜ r L ≥ 1 will be replaced by the slightly stronger constraint b L +1 + ˜ r ≥ 1 .

E.1.4 t -th forward pass

Formally, we sum the updates in each step,

where δW L +1 t +1 := -η ˜ χ t (˜ x L t ) T denotes the normalized change in the weights W L +1 (as a row vector) of scaling θ L +1 = θ W L +1 = n -c L +1 under perturbation stability and nontriviality so that ˆ W L +1 t scales as θ ∇ = n -min( b L +1 ,c L +1 ) . δW L +1 t +1 should not be confused with ˜ δW L +1 t +1 which denotes the perturbation of the weights at time t +1 . For every nontrivial stable parametrization we have ˜ χ t = Θ(1) and ˜ x L t = Θ(1) which requires ˜ θ L ≤ 1 . In the case c L +1 < b L +1 , we write ˆ W L +1 t := n -b L +1 + c L +1 ˆ W L +1 0 + ( δW L +1 1 + · · · + δW L +1 t ) with the same scaling θ ∇ = n -min( b L +1 ,c L +1 ) .

For preactivations and activations we also sum the changes from each step,

Using the fact that yields the normalized preactivation updates

with scaling θ 1 = θ W 1 = n -c 1 θ SAM = n -c 1 -min( b L +1 ,c L +1 ,d + d L +1 ) as for SGD under perturbation stability and nontriviality where ˜ χ t -1 = Θ(1) .

For l ∈ [2 , L ] , it holds that

with the right scaling θ W l = n 1 -c l -min( b L +1 ,c L +1 ,d + d L +1 ) as for SGD under perturbation stability ˜ x l -1 t -1 = Θ(1) , so that we get δh l t using a telescope sum,

with scaling θ ∇ = n -min( b L +1 ,c L +1 ) .

As in the first backward pass, we use the weight updates to directly compute the preactivation and activation perturbations similar to the t -th forward pass but performing SGD instead of SAM in the last step. The SGD backward pass for the perturbation is given by

with scaling max( θ ∇ , n 1 -c l θ SAM θ ∇ ) = θ ∇ , since n 1 -c l θ SAM ≤ 1 is implied by r ≥ 0 required for the stability of (pre-)activation updates.

We write χ t = L ′ ( f t ( ξ t ) , y t ) for the derivative of the loss with respect to the unperturbed function (which is Θ(1) under stability and nontriviality), and get

which defines a NonLin operation with the vectors W l 0 δx l -1 t , dh l SAM, 0 , dh l SAM,t -1 and everything else treated as scalars. The scaling is given by

with

̸

where θ W l ≤ 1 for all l ∈ [ L ] for stability. Note that for l 1 ≤ l 2 , it holds that θ l 1 ≤ θ l 2 , which explains the sufficiency of θ L = n -r L = n -r for the stability of the activation updates.

Activations with the same scaling θ l can then simply be defined via the NonLin operation

The updates of the output function are scalars defined as

where θ ′ L +1 = nθ L +1 = n 1 -c L +1 and θ ′ L ∇ = nθ ∇ θ L = n 1 -min( b L +1 ,c L +1 ) -r L , where we will see why W L +1 t -1 = Θ( n -min( b L +1 ,c L +1 ) ) in the next paragraph. This leads to the constraints c L +1 ≥ 1 and b L +1 + r ≥ 1 for the stability of the output function, where equality in either constraint leads to nontriviality.

E.1.5 t -th backward pass

Perturbations. Due to linearity and stability, the last layer remains

which defines a NonLin operation with the vectors W l 0 ˜ δx l -1 t , dh l SAM, 0 , . . . , dh l SAM,t -1 , dh l t , and where we can now define the definitive scalings ˜ θ 1 := ˜ θ W 1 := n -( d + d 1 ) θ ∇ = n -(min( b L +1 ,c L +1 )+ d + d 1 ) , ˜ θ W l := n 1 -( d + d l ) θ ∇ = n -(min( b L +1 ,c L +1 )+ d +( d l -1)) and ˜ θ l = max( ˜ θ l -1 , n 1 -c l θ SAM ˜ θ l -1 , ˜ θ W l ) = max l m =1 ˜ θ W m = n -˜ r l with

̸

where we used that n 1 -c l θ SAM ≤ 1 due to r ≥ 0 for stability and ˜ x l -1 t = Θ(1) due to perturbation stability. Perturbation stability of the hidden layer (pre-)activations ˜ δh l , ˜ δx l = O (1) for all l ∈ [ L ] holds if and only if ˜ r := ˜ r L ≥ 0 since ˜ r l ≥ ˜ r L for all l ≤ L .

The activation perturbations ˜ δx l t and the perturbation of the output function ˜ δf t can be defined exactly as in the first backward pass,

with ˜ δW L +1 t := ρ χ t x L t ∥ v t ∥ and the same scalings ˜ θ l , ˜ θ ′ L +1 = n 1 -( d + d L +1 ) and ˜ θ ′ L ∇ = nθ ∇ ˜ θ L = n 1 -min( b L +1 ,c L +1 ) -˜ r since W L +1 t = W L +1 0 +∆ W L +1 t = max( n -b L +1 , n -c L +1 ) , which yields the slightly stronger constraint (than in the first backward pass) min( b L +1 , c L +1 )+˜ r ≥ 1 for perturbation stability and either ˜ θ ′ L +1 = 1 or min( b L +1 , c L +1 ) + ˜ r = 1 for perturbation nontriviality.

SAM Update. For each l ∈ { 1 , . . . , L } , as in the first backward pass, we get

with scaling θ SAM = n -min( b L +1 ,c L +1 ,d L +1 +1 / 2) as well as

For dx l SAM,t we again use a telescope sum over the weight changes,

which defines a NonLin operation in the vectors ( W l 0 ) T dh l SAM,t , ˜ x l -1 0 , . . . , ˜ x l -1 t -1 , x l -1 t and everything else treated as scalars. Note that the scalings remain θ SAM , since ∇ x l -1 t f | ˜ W t = Θ(max( θ SAM , n 1 -c l θ 2 SAM , n 1 / 2 -d l θ ∇ θ SAM )) = Θ( θ SAM ) under stability, nontriviality, perturbation stability and perturbation nontriviality.

Finally define the loss derivative on the perturbed output function

and compute the normalized change in W L +1 ,

E.2 The infinite-width limit

In this section, we apply the Master Theorem's computation rules to derive the marginal distributions Z corresponding to the vectors of the program constructed above. According to the Master Theorem, each such vector z will have roughly iid coordinates distributed like Z z in the large n limit.

We assume stability holds, so that θ → ˚ θ ∈ { 0 , 1 } for all scalars θ in the program.

For the first forward pass, we have

If b L +1 > 1 / 2 then ˚ f 0 = 0 , otherwise if b L +1 = 1 / 2 then ˚ f 0 converges to a nontrivial Gaussian. For the details we refer to Appendix H.4.1 in Yang and Hu (2021), as at initialization their results still hold here.

For the first SGD backward pass, we have

Z dx L 0 ( ξ ) = Z ˆ W L +1 0 , Z dh l 0 ( ξ ) = Z dx l 0 ( ξ ) ϕ ′ ( Z h l 0 ( ξ ) ) , Z dx l -1 0 ( ξ ) = Z ( W l 0 ) T dh l 0 ( ξ ) , where ˙ Z dx l 0 ( ξ ) = 0 and Z dx l 0 ( ξ ) = ˆ Z dx l 0 ( ξ ) for all ξ ∈ X .

For general t > 0 , we have Z dx L t ( ξ ) = Z ˆ W L +1 t , Z dh l t ( ξ ) = Z dx l t ( ξ ) ϕ ′ ( Z h l t ( ξ ) ) , Z dx l -1 t ( ξ ) = Z ( W l 0 ) T dh l t ( ξ ) -η ˚ θ W l t ∑ s =1 ˚ ˜ χ s -1 E [ Z dh l SAM,s -1 Z dh l t ] Z ˜ x l -1 s -1 , where ˚ ˜ χ s = L ′ ( ˚ ˜ f s ( ξ s ) , y s ) for s < t , and Z ( W l 0 ) T dh l t ( ξ ) is a Θ(1) random variable distributed as

For all t ≥ 0 , the limit of the gradient norm is given by

where ˚ χ t = L ′ ( ˚ f t ( ξ t ) , y t ) , θ 2 ∥ v 1 ∥ := n 1 -2 d 1 θ 2 ∇ , θ 2 ∥ v l ∥ := n 2 -2 d l θ 2 ∇ for l ∈ [2 , L ] and θ 2 ∥ v L +1 ∥ := n 1 -2 d L +1 , and where ˚ θ 2 ∥ v L +1 ∥ = 1 if and only if d L +1 = 1 / 2 and ˚ θ 2 ∥ v L +1 ∥ = 0 if and only if d L +1 > 1 / 2 , while ˚ θ 2 ∥ v l ∥ = 1 if and only if 2 d l = 1 + I ( l > 1) -2 min( b L +1 , c L +1 ) and ˚ θ 2 ∥ v l ∥ = 0 if and only if 2 d l > 1 + I ( l > 1) -2 min( b L +1 , c L +1 ) .

For the last-layer weight perturbations (for θ ∇ ≥ ˜ θ L +1 , else Z ˆ ˜ W L +1 t = Z ˜ δW L +1 t ) we have

Note that ˚ χ t cancels itself out and we purely get a perturbation in distribution Z x L t scaled to have standard deviation ρ .

For all t ≥ 0 and l ∈ [1 , L ] , we have

where for l = 1 ,

If ˚ ˜ θ l = 0 , then otherwise ˚ ˜ θ l = 1 and

For l ≥ 2 , we have

where ˜ θ ( l -1) /l = ˜ θ l -1 ˜ θ l , θ W l ( ˜ l -1) / ˜ l = θ Wl ˜ θ l -1 ˜ θ l and ˜ θ W l /l = ˜ θ Wl ˜ θ l , and Z W l 0 ˜ δx l -1 t has the decomposition

The perturbed output function has the limit ˚ ˜ f t := ˚ f t + ˚ ˜ δf t with

so that we can define ˚ ˜ χ t = L ′ ( ˚ ˜ f t ( ξ t ) , y t ) or equivalently ˚ ˜ χ t = L ′ ( ˚ ˜ θ L +1 ˚ ˜ θ L E [ Z ˆ ˜ W L +1 t Z ˜ x L t ] , y t ) . For the SAM gradients we have

where Z ( W l 0 ) T dh l SAM,t is given by

Now SAM's (pre-)activation updates are given by

where θ ( l -1) /l := θ l -1 /θ l , θ W l /l := θ W l /θ l and Z W l 0 δx l -1 t has the decomposition

otherwise ˚ θ l = 1 and

The last-layer SAM weight update is given by

For t > 0 , the SAM function update is given by

E.3 Concluding the proof of all main results

After writing out the NE ⊗ OR ⊤ 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) in the following way.

̸

Formally Yang and Hu (2021) show feature learning for SGD with small enough learning rate η > 0 by proving ∂ 2 η E ( Z x L 1 ( ξ 0 ) ) 2 = 0 at η = 0 , and they show that learning does not occur in the kernel regime by showing ∂ 3 η ˚ f 1 = 0 , hence ˚ f 1 -˚ f 0 is not linear in η .

̸

Both E ( Z x L 1 ( ξ 0 ) ) 2 and ˚ f 1 are defined via NE ⊗ OR ⊤ computations and can be written as a composition of additions, multiplications, the expectation operator, applications of ϕ and ϕ ′ , overall applications of infinitely differentiable, pseudo-Lipschitz functions to (Gaussian) random variables, η and ρ . Consequently E ( Z x L 1 ( ξ 0 ) ) 2 and ˚ f 1 are infinitely often differentiable as a function of both η and ρ , where differentiating the expectation operator is covered in Yang and Hu (2021, Lemma H.39). Since Yang and Hu (2021) cover the case ρ = 0 , their proofs immediately show the correctness of the derived scalings for SAM as long as η > 0 and ρ > 0 are chosen small enough. Both the gradient evaluation for the perturbation as well as the gradient evaluation for the updates stay arbitrarily close to those of SGD if ρ > 0 is chosen small enough. The conditions for stability, nontriviality, feature learning, perturbation nontriviality and effective perturbations now follow from considering the respective scaling.

E.3.1 Proof of Theorem D.2

A bcd -parameterization is stable if and only if all scalings in the Tensor Program have the limit ˚ θ ∈ { 0 , 1 } , where ˚ θ = 1 is required for activations at initialization (for which nothing changes compared to SGD). Potential cancellations are taken care of for sufficiently small η > 0 and ρ > 0 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

A stable bcd -parameterization is nontrivial if and only if ˚ f t = Θ(1) if and only if ˚ θ ′ L +1 = 1 or ˚ θ ′ L ∇ = 1 .

E.3.3 Proof of Theorem D.4

A stable bcd -parametrization is feature learning in layer l if and only if the feature update scaling ˚ θ l = 1 where

̸

Hence a stable bcd -parametrization is feature learning in layer l if and only if r l = 0 .

Since for all l 1 ≤ l 2 , it holds that r l 1 ≥ r l 2 ≥ 0 , we get the equivalence for any l 0 ∈ [ L ] : A stable bcd -parametrization is feature learning in layer l 0 if and only if it is feature learning in layer l for all l ≥ l 0 if and only if r l 0 = 0 .

E.3.4 Proof of Theorem D.6

Given a stable bcd -parametrization, perturbation triviality is fulfilled if and only if ˚ ˜ θ ′ L +1 = 0 and ˚ ˜ θ ′ L ∇ = 0 , where ˜ θ ′ L +1 = n 1 / 2 -d L +1 and ˜ θ ′ L ∇ = nθ ∇ ˜ θ L = n 1 -min( b L +1 ,c L +1 ) -˜ r , hence if and only if d L +1 > 1 / 2 and min( b L +1 , c L +1 ) + ˜ r > 1 .

̸

In that case, ˚ ˜ f t = ˚ f t , but ˚ f t may still be affected by non-vanishing SAM perturbations in δW L +1 t and δx L t . Only when all SAM perturbations vanish are we effectively only using SGD. By definition, the perturbation scale in the l -th layer vanishes if and only if ˚ ˜ θ l = 0 , where ˜ θ l = n -˜ r l with ˜ r l = min( b L +1 , c L +1 ) + 1 / 2 + min l m =1 ( d m -I ( m = 1)) , hence if and only if ˜ r l > 0 . Since ˜ r l ≥ ˜ r L = ˜ r for all l ≤ L , we get ˚ ˜ θ l = 0 for all l ∈ [ L ] if and only if ˜ r > 0 . Similarly, for any reference layer l 0 ∈ [ L ] , we get ˚ ˜ θ l = 0 for all l ≤ l 0 if and only if ˜ r l 0 > 0 . In words, for any l 0 ∈ [ L ] , we have vanishing perturbations in layer l 0 if and only if we have vanishing perturbations until layer l 0 if and only if ˜ r l 0 > 0 .

Altogether, a stable bcd -parametrization has vanishing perturbations if and only if ˜ r > 0 , d L +1 > 1 / 2 and min( b L +1 , c L +1 ) + ˜ r > 1 . This case reduces to the results in Yang and Hu (2021) in the limit. Since stability requires c L +1 ≥ 1 and ˜ r ≥ 0 , we can rewrite the equivalence conditions as d L +1 ≥ 1 / 2 and ˜ r > max(0 , 1 -b L +1 ) .

E.3.5 Proof of Theorem D.8

Recall ˜ θ W 1 := n -( d + d 1 ) θ ∇ , ˜ θ W l := n 1 -( d + d l ) θ ∇ and, for the last layer ˜ θ W L +1 := n -( d + d L +1 ) .

As opposed to perturbation nontriviality, we are not only interested in ˜ θ l = max( ˜ θ l -1 , ˜ θ W l ) = max l m =1 ˜ θ W m → 1 , but in a non-vanishing contribution of the perturbations in layer l , i.e. ˚ ˜ θ W l = 1 or, for the last layer, ˚ ˜ θ L +1 = 1 .

E.3.6 Proof of Theorem D.9

The limit of the gradient norm is defined as a NE ⊗ OR ⊤ program scalar (E.3). Note that for b L +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 { b l } l ∈ [ L +1] and learning rates { c l } l ∈ [ L +1] . To fulfill the gradient norm constraints (D.1), we have to choose d l = C = 1 / 2 for all l ∈ [ L +1] , because stability requires min( b L +1 , c L +1 ) ≥ 1 / 2 . Now stability of the output 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 ˜ δW L +1 t ˜ x L t . After choosing d ≥ 1 / 2 , we get ˜ r ≥ min( b L +1 , c L +1 ) ≥ 1 / 2 > 0 which implies vanishing perturbations in all hidden layers.

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 L +1 = 1 / 2 , which requires d ≥ 1 / 2 for stability of the output function perturbations. Achieving non-vanishing gradient norm contributions of all layers requires d 1 = 1 / 2 -min( b L +1 , c L +1 ) and d l = 1 -min( b L +1 , c L +1 ) for l ∈ [2 , L ] , which results in ˜ r = d ≥ 1 / 2 > 0 which implies vanishing perturbations in all hidden layers.

E.3.8 Proof of Theorem D.11

Given a stable bcd -parametrization, we know d + d L +1 ≥ 1 , so that the feature learning constraint r is not affected by any stable choice of d ∪ { d l } l ∈ [ L +1] . The maximal stable choice of layerwise initialization variances { b l } l ∈ [ L +1] and learning rates { c l } l ∈ [ L +1] that constitute µ P is therefore unaffected by the perturbation scalings d ∪ { d l } l ∈ [ L +1] .

Stability of the output function perturbations requires b L +1 + ˜ r ≥ 1 . Hence if b L +1 < 1 , then ˜ r ≥ 1 -b L +1 > 0 , which implies vanishing perturbations in all hidden layers.

̸

To show uniqueness we iterate through all possibilities of saturating the norm bound constraint (D.1). We have considered the cases d L +1 = 1 / 2 in (b) leading to vanishing perturbations in all hidden layers and d 1 = 1 / 2 -c ∇ in (c) with only one choice for effective perturbations in all layers. Lastly consider d l = 1 -c ∇ for l ∈ [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 l ∈ [2 , L ] requires ˜ θ W l = 1 for which we need d = 0 . But then, as d 1 ≥ 1 / 2 -c ∇ , it holds that ˜ r 1 ≥ 1 / 2 implying vanishing perturbations in the first layer. This shows the uniqueness of (1).

From now on consider b L +1 ≥ 1 . Recall c ∇ := min( b L +1 , c L +1 ) . In µ P, c ∇ = 1 , but effective perturbations in all layers can be achieved more generally for c ∇ ≥ 1 . Choosing d 1 = 1 / 2 -c ∇ saturates the gradient norm constraint (D.1). To reach effective perturbations already in the first layer ˜ r 1 = c ∇ + d + d 1 = 0 , we need d = -1 / 2 . For perturbation stability and last-layer effective perturbations, we need d + d L +1 = 1 which requires d L +1 = 3 / 2 . Achieving perturbation stability and effective perturbations in all hidden layers requires ˜ θ W l = 1 which is equivalent to c ∇ + d + d l -I ( l = 1) = 0 . For l ∈ [2 , L ] , we therefore need d l = 3 / 2 -c ∇ . This choice of { d l } l ∈ [ L +1] achieves effective perturbations in all layers.

For the gradient norm statements, note that the gradient norm ∥ v t ∥ can be written as a NE ⊗ OR ⊤ computation rule (E.2) where the layer scalings in this parameterization are Θ(1) for the input layer, Θ( n -1 / 2 ) 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 is equivalent to ˜ r = 0 . Since min( b L +1 , c L +1 ) ≤ 1 , we get min( b L +1 , c L +1 )+˜ r ≤ 1 . Since stability requires min( b L +1 , c L +1 )+ ˜ r ≥ 1 , we get min( b L +1 , c L +1 ) + ˜ r = 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 ˜ r l ≤ ˜ r 1 = 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 a closer analysis of its fine-grained dynamics in comparison to SGD to future work. Before looking into the effective perturbation regime, we restate Lemma H.37 in Yang and Hu (2021) with a more detailed proof.

First, we define ℓ ∈ [ L ] as the unique index that satisfies θ L = · · · = θ ℓ = 1 > θ ℓ -1 ≥ · · · ≥ θ 1 . In words, ℓ is the first layer in which feature learning occurs. Analogously, we define ˜ ℓ ∈ [ L ] as the unique index that satisfies 1 = ˜ θ L ˜ θ L = · · · = ˜ θ ˜ ℓ ˜ θ L > ˜ θ ˜ ℓ -1 ˜ θ L ≥ · · · ≥ ˜ θ 1 ˜ θ L .

Lemma E.2 ( Features after first SGD step ) . Defining Z l t := Z h l t , γ l ( η ) = E ϕ ( Z l 0 ) ϕ ( Z l 1 ) for l ≥ 1 , γ 0 = ξ T 0 ξ and γ l 11 ( η ) = E ϕ ′ ( Z l 0 ) ϕ ′ ( Z l 1 ) , we have

and, for all l ≥ ℓ , where β l is defined recursively by

with β ℓ -1 = 0 . Note that β l (0) < 0 for all l ≥ ℓ .

Proof. By the defining infinite-width equations, assuming ˚ θ W l /l = 1 (so minimal stable choice of c l ),

At l = ℓ , we get ˚ θ ( ℓ -1) /ℓ = 0 , whereas for l > ℓ we get ˚ θ ( l -1) /l = 1 , which results in ˚ θ ( ℓ -1) /ℓ = I l>ℓ . Now, for l > ℓ , the second term decomposes into ˆ Z W l 0 δx l -1 1 and

Since by induction hypothesis,

where Z dx l -1 0 = Z ( W l 0 ) T dh l 0 is the only dependence on ˆ Z ( W l 0 ) T dh l 0 , we get

Plugging the derivative back into the defining equation and noticing that Z dh l 0 = Z dx l 0 ϕ ′ ( Z l 0 ) concludes the proof.

An analogous analysis for the perturbation at initialization shows.

Lemma E.3 ( Feature perturbation at initialization ) . The perturbation trivial layers fulfill

and, for all l ≥ ˜ ℓ , where ˜ β l independent of η is defined recursively by

with ˜ β ˜ ℓ -1 = 0 . Note that ˜ β l (0) > 0 for all l ≥ ˜ ℓ .

)

ϕ

(

Z

h

l

0

-

)]

by

ξ

0

T

ξ

.

If

ℓ

˜

= 1

, in the definition of

β

˜

l

replace

E

[

ϕ

(

Z

h

l

0

-

Now we are ready to state the closed form expression for the first SAM update.

Lemma E.5 ( Features after first SAM update ) . Defining Z l t := Z h l t and ˜ Z l t := Z ˜ h l t , we have

1

˜

1

and, for all l ≥ ℓ ,

where β l is defined recursively by

with β ℓ -1 = 0 , and γ l = γ l ( η ) is recursively defined by

Remark E.6.

If

ℓ

= 1

, in the definition of

β

l

replace

E

[

ϕ

( ˜

Z

l

-

0

1

)

ϕ

(

Z

l

-

1

1

)]

by

(

ξ

T

ξ

)

.

t

-

1

◀

◀

Proof. By the defining infinite-width equations, for l ≥ ℓ , assuming ˚ θ W l /l = 1 (so minimal stable choice of c l ),

At l = ℓ , we get ˚ θ ( ℓ -1) /ℓ = 0 and ˚ θ W ℓ /ℓ = 1 , whereas for l > ℓ we get ˚ θ ( l -1) /l = 1 and ˚ θ W l /l = 1 (under minimal stable choice of c l ), which results in ˚ θ ( l -1) /l = I l>ℓ . Now, for l > ℓ , the second term decomposes into ˆ Z W l 0 δx l -1 1 and ˙ Z W l 0 δx l -1 1 . For the rest of the proof it remains to analyse ˙ Z W l 0 δx l -1 1 . Since by induction hypothesis,

+

where Z dx l -1 SAM, 0 = Z ( W l 0 ) T dh l SAM, 0 + ρ ˚ ˜ θ W l ˚ χ 0 ˚ ¯ ∥∇ L 0 ∥ E [ Z dh l 0 Z dh l SAM, 0 ] Z x l -1 0 with the second term independent of ( W l 0 ) T and by Lemma E.3 we know ˜ Z l -1 0 = Z l -1 0 + I l -1 > ˜ ℓ ˆ Z W l -1 0 ˜ δx l -2 0 ρ ˜ β l -1 Z dx l -1 0 ϕ ′ ( Z l -1 0 ) , where only the last term with Z dx l -1 0 = Z ( W l 0 ) T dh l 0 influences ˙ Z W l 0 δx l -1 1 , we get

with

and, using Z dh l -1 0 = Z dx l -1 0 ϕ ′ ( Z l -1 ) = Z ( W l 0 ) T dh l 0 ϕ ′ ( Z l -1 ) , yields

Plugging Eq. (E.5) back into the defining equation (E.4) and noticing that Z dh l SAM, 0 = Z dx l SAM, 0 ϕ ′ ( ˜ Z l 0 ) as well as Z dh l 0 = Z dx l 0 ϕ ′ ( Z l 0 ) concludes the proof.

F Generalizations and further perturbation scaling considerations

F.1 Overview over choices of d l and d

Since for some combinations of architectures and datasets it turns out that performing SAM on a 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 d l 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 d l from µ P 2 for the perturbed layers. This section is instead interested in a complete characterization of all possible choices of { d l } l ∈ [ L +1] and d . The heuristic derivation only requires the gradient norm constraints (D.1) and the perturbation stability constraints that require ˜ δW 1 = O (1) and ˜ δW l = O ( n -1 ) for l > 1 given by

where a layer is 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

perturbation scalings d and d l in Table F.1 and allows us to formulate a simple rule of how to choose d and d l 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 L SAM ⊆ L .

If all input-like layers should have vanishing weight perturbations but there exists a hiddenlike layer l ∈ L SAM , set d = 0 . Hidden-like layers are effectively perturbed if and only if d l = 1 -c ∇ . Output-like layers are effectively perturbed if and only if d L +1 = 1 . For all layers that have vanishing weight perturbations, do not perturb these weights, or set d l > c ∇ for input-like, d l > 1 -c ∇ for hidden-like and d l > 1 for output-like layers (as required by the perturbation stability and gradient norm constraints).

If there exists an input-like layer l ∈ L SAM , set d = -1 / 2 . Input-like layers are effectively perturbed if and only if d l = 1 / 2 -c ∇ . Hidden-like (respectively, output-like) layers are effectively perturbed if and only if d l = 3 / 2 -c ∇ (respectively, d l = 3 / 2 ). For all layers that have vanishing weight perturbations, do not perturb these weights or choose d l > 1 / 2 -c ∇ for input-like, d l > 3 / 2 -c ∇ for hidden-like and d l > 3 / 2 for output-like layers.

If both all input-like and all hidden-like layers have vanishing weight perturbations, but there exists some output-like layer l ∈ L SAM , then set d = 1 / 2 . Output-like layers are effectively perturbed if and only if d l = 1 / 2 . For all layers that have vanishing weight perturbations, do not perturb these weights or set d l ≥ 1 / 2 -c ∇ for input-like, d l ≥ 1 -c ∇ for hidden-like and d l > 1 / 2 for output-like layers (as required by the perturbation stability and gradient norm constraints).

Example F.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.1). To achieve first-layer effective perturbations, but trivial weight perturbations in all other layers, we need ˜ θ W 1 = 1 and ˚ ˜ θ W l = 0 , for which we will choose ˜ θ W l = n -1 . This requires setting

If L SAM = ∅ , then set d > 1 / 2 or simply perform SGD.

where one of the constraints (D.1) has to be fulfilled. Plugging the above d l -choices into (D.1) results in the constraints d ≤ -1 / 2 , d ≤ 1 , d ≤ 3 / 2 , hence choose d = -1 / 2 so that only the first layer contributes non-vanishingly to the gradient norm. Note that ˜ r = 0 and output perturbation nontriviality holds if and only if min( b L +1 , c L +1 ) = 1 (as in µ P). We apply this perturbation scaling in Appendix H.2 to show that propagating perturbations from early layers are not enough to inherit SAM's inductive bias that leads to improved generalization performance. ◀

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 adapted optimization algorithm should recover the original (SAM) algorithm when not using layerwise perturbation scaling.

The perturbation step should require at most one additional forward and backward pass in each update step.

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

where ∥ · ∥ 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 ρ l ! = Θ( √ fan_out / fan_in ) . ◀

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); Andriushchenko 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 in this paper is to 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 to Definition 4, for this variant the perturbation norm, i.e. the radius of the adversarial ascent ball, is not guaranteed to be ρn -d , but ρ ( ∑ l ∈ [ L +1] ρ 2 l ) 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

the correct perturbation scaling in µ P is given by ρ l ! = Θ( n 1 / 2 · fan_out / fan_in ) .

To understand this scaling rule, note that for b L +1 > 1 / 2 (such as in µ P), the last layer always dominates the gradient norm (see Eq. (E.2) for the TP argument), resulting in the scaling

Thus, compared to SAM without gradient normalization (Appendix F.3), ∥∇ W L∥ F always contributes the scaling n 1 / 2 . Noting that perturbations should scale like updates, and updates receive the layerwise learning rates η l ! = Θ( fan_out / fan_in ) concludes the derivation. ◀

In Definition 4, we 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 ρn -d . 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 be scaled like updates.

Another alternative to layerwise perturbation scaling as in Definition 4 is motivated by the observation, that in µ P 2 with Definition 4, only the first layer dominates the joint gradient norm (Theorem D.11). To let all layers contribute width-independently to the joint gradient norm, we can introduce even more hyperparameters (with limited benefit) by decoupling the numerator and denominator scalings in the perturbation. Opposed to Definition 4, the perturbation norm is again not analytically set with the choice of ρ , 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 1 for details).

Remark F.4 ( SAM with decoupled perturbation numerator and denominator scaling ) . For (SAM) with perturbations

ρ n -d l n -d l ◀

with layerwise perturbation radii · and separate gradient norm scaling ˜ .

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 a layerwise gradient rescaling can also be achieved by adapting the architecture with widthdependent weight multipliers. The multipliers ( a -µP 2 ) achieve effective perturbations without layerwise perturbation scaling such that all layers contribute non-vanishingly 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-based analysis. Our TP-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 ∥∇ W L∥ F , we just drop all perturbation scaling terms n -d l . For the examples of (LN) and (DP), we replace (E.2) in each layer separately by the normalized TP scalars,

with scaling θ ∥∇ 1 ∥ = n 1 / 2 θ ∇ for the first layer, where θ ∇ is overloaded to denote θ ∇ = n -b L +1 in the first step and θ ∇ = n -min( b L +1 ,c L +1 ) in all later steps (in µ P, θ ∇ = n -1 always),

with scaling θ ∥∇ L +1 ∥ = nθ ∇ for all hidden layers l ∈ [2 , L ] , and

with scaling √ n for the output layer, with respective normalized limits

where ˚ χ t = L ′ ( ˚ f t ( ξ t ) , y t ) .

The adapted scalings can then be tracked as before to derive the maximal stable layerwise perturbation scaling. Consider for example input layers in (LN). In µ P, we know ∥∇ W 1 L t ∥ = Θ( n -1 / 2 ) and ∇ W 1 L t = Θ( n -1 ) entrywise. Effective perturbations are achieved with ε 1 = Θ(1) , so choose ρ l = n 1 / 2 as expected from ( ∗ ). Proceed similarly for all layers and perturbation scaling variants.

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

The NE ⊗ OR ⊤ program for this update rule with arbitrary v l t = n -c l ∇ w L ( f ( ξ t ; W t ) is also easily adapted from the above derivation. Just note that the gradient norm appears in an equation if and only if the perturbation radius ρn -d appears. Without dividing by ∥ v t ∥ , the parameter d becomes superfluous. Simply set d = 0 and remove the gradient norm constraints (D.1) to arrive at the NE ⊗ OR ⊤ program and bcd -constraints for the update rule without gradient normalization.

Perturbation scaling d l plays a similar role as learning rate scaling c l as both scale similar gradients. We get effective perturbations in the l -th layer from the equation d l + min( b L +1 , c L +1 ) = c l +

min( b L +1 , c L +1 , d L +1 ) in µ P, which yields d l = c l for all l ∈ [ L ] (since d L +1 = 1 for stability). In particular, in µ P, the correct layerwise perturbation scaling of unnormalized gradients is given by the rule fan out fan in or the squared weight (update) spectral norm ∥ W l ∥ 2 ∗ (Yang et al., 2023a), which could be efficiently approximately tracked with a running power iteration.

Note that Dai et al. (2024) argue that the normalizing the gradients for the perturbation is crucial (in standard parametrization) due to a 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 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 to scale down ρ less aggressively in practice.

F.4 Extension to Adaptive SAM

Adaptive SAM (ASAM) (Kwon et al., 2021) is motivated by a 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.5. Neither elementwise ASAM, which performs (SAM) but using the perturbation rule (F.4) , nor layerwise ASAM, which performs (SAM) but using the perturbation rule (F.6) , can be written as a NE ⊗ OR ⊤ program.

Proof sketch. Elementwise ASAM requires an elementwise multiplication of matrices, and layerwise ASAMrequires calculating the Frobenius norm of a matrix. A NonLin operation only takes vectors as arguments, so NE ⊗ OR ⊤ calculations with a matrix require its multiplication with a 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 NE ⊗ OR ⊤ master theorem states that coordinates of NE ⊗ OR ⊤ vectors behave iid-like, this mixture cannot be disentangled by choosing a structured vector. Hence, already at initialization, the square of individual entries/the Frobenius norm of a random matrix cannot be exactly recovered by a function of matrix-vector products with NE ⊗ OR ⊤ vectors.

Although ASAM is not formally covered by our theory, we still expect that the ASAM perturbations are correlated with the gradient and therefore with the incoming activations, so that heuristically we can still expect LLN-like behaviour and apply our scaling condition. If the perturbation rules still behave LLN-like, then Table 1 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:

In µ P, it always holds that

Assuming the normalization term in the denominator is scaled to Θ(1) , track the layerwise scalings of the numerator. Maximal stable perturbations are always achieved with

This yields constraints for achieving maximal stable perturbations in each layer.

To ensure normalization term scaling Θ(1) , iterate through the layers l :

Now replace the norm constraints (D.1) by tracking the scalings of each layer's contribution to the update rule's total normalization term.

(a) choose d l to satisfy its norm constraint,

(b) choose d to induce maximal stable perturbations in that layer,

̸

From the above configurations, choose the unique one that yields maximal stable perturbations in all layers. 3

(c) choose all other d l ′ , l ′ = l , minimal to both satisfy its norm constraint as well as perturbation stability ˜ δW l t = { O (1) l = 1 , O ( n -1 ) l > 1 .

F.4.1 Elementwise ASAM

If we want to be invariant to elementwise rescaling operators T l w ( x ) = | W l | ⊙ x where x, W l ∈ R m × n and ⊙ denotes elementwise multiplication, the resulting ASAM perturbation rule (where we introduce (layer-wise) perturbation scalings { d } ∪ { d l } l ∈ [ L +1] ) replaces (LP) and is given by with normalization

where the absolute values | W l | 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 µ P. In µ P, the layerwise weights and gradients scale as (F.2). For ∥∇ elem ASAM ∥ = O (1) , we therefore replace the constraints (D.1) by the constraints where we can choose { d l } l ∈ [ L +1] to achieve equality in at least one constraint to achieve ∥∇ elem ASAM ∥ = Θ(1) .

The layerwise perturbations scale as ˜ δW l t = n -d    Θ( n -d 1 θ ∇ ) l = 1 , Θ( n -1 -d l θ ∇ ) l ∈ [2 , L ] , Θ( n -d L +1 n -2 ) l = L +1 .

Stable and nontrivial perturbations in each layer are achieved under condition (F.3), which induces the constraints for optimal layerwise perturbation scaling

Irrespective which of the above norm constraints (F.5) we satisfy, we need d = -1 / 2 to achieve optimal layerwise perturbation scaling. Hence d = d L +1 = -1 / 2 and d l = 1 / 2 -c ∇ for l ∈ [ L ] is the unique choice of { d } ∪ { d l } l ∈ [ L +1] modulo norm scaling equivalence that achieves Θ(1) perturbation scaling in all layers. With this choice all layers contribute non-vanishingly to the gradient norm. In µ P c ∇ = 1 , so that d l = -1 / 2 for all l ∈ [ L + 1] , so that ASAM does not require layerwise rescaling of the gradients, but upscaling of the perturbation by n 1 / 2 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 if the perturbation radius is tuned globally at every width.

If instead of a global gradient norm ∥∇ elem ASAM ∥ , one would want to normalize in each layer separately with ∥∇ elem,l ASAM ∥ := n -d l ∥| W l | ⊙ ∇ W l L ( f ( ξ t ; W t ) , y t ) ∥ F , the layerwise perturbation scalings become ˜ δW l t = n -d Θ( n -3 / 2 ) l > 1 . Again, to achieve maximal stable perturbations in all

d = -1 / 2

{ Θ( n -1 / 2 ) l = 1 , layers we need and no layerwise adaptation of the gradient norm.

F.4.2 Layerwise ASAM

ASAM with layerwise rescaling as in Müller et al. (2024) employs the layerwise transformations T l w ( x ) = ∥ W l ∥ F · x . This ASAM perturbation rule replaces (LP) and is given by with normalization

Hence, the norm constraints (D.1) are now replaced by

The scale of the perturbation numerator now scales as ˜ δW l t = n -d    Θ( n -d 1 nθ ∇ ) l = 1 , Θ( n -d l nθ ∇ ) l ∈ [2 , L ] , Θ( n -d L +1 n -1 ) l = L +1 .

In µ P, achieving maximal stable perturbations (F.3) is therefore equivalent to satisfying the constraints

Now we can simultaneously satisfy the first- and last-layer norm constraints with d 1 = 0 and d L +1 = 0 , while achieving effective perturbations in all layers with d = 0 and d l = 1 . Satisfying the norm constraint in the hidden layers with d l = 1 / 2 would imply vanishing perturbations in the first and last layer (by requiring d ≥ 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 in ResNets and ViTs in 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 normalization ∥ v t ∥ 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 as

where ε > 0 is a small positive constant, γ l t , β l t are learnable parameters and ν l t = 1 n ∑ n i =1 ( x l t ) i is an Avg operation as in Yang and Littwin (2023, Def. 2.6.1) and σ l t = √ 1 n ∑ n i =1 ( x l t -ν l t ) 2 is a composition of Nonlin, Avg and Nonlin. The parameters γ l t , β l t can be seen as input weights to the input 1 . They should be initialized as γ l 0 = 1 and β l 0 = 0 . In the forward pass, the layernorm preserves stability h l +1 t = Θ(1) when γ l t + β l t = Θ(1) except for the Lebesgue nullset of learning rates for which they exactly cancel each other out. Recall the notation dz = θ -1 z ∂f/∂z , where θ z = n C for some C ∈ R denotes the width-dependent scaling. The derivatives are

These gradients coincide both in shape and scaling with the scaling we expect for an input layer, resulting in the same gradient spectral/Frobenius norm scaling. Continuing the backward pass, using

which preserves the order as long as γ l t = Θ(1) , since x l t = Θ(1) , we know ν l t , σ l t = Θ(1) .

Note that Layernorm removes the necessity to avoid blowup in the activations x l t in the forward pass (ignoring potential numerical issues), and always rescales to Θ(max( γ l t , β l t )) . However, in the backward pass, a scaling x l t = Θ( n c ) , with c > 0 , results in dx l t = Θ( n -c dh l +1 t γ l t ) , hence vanishing gradients. The gradients would only stabilize if ϕ ′ ( h l t ) = Θ( h l t ) , but no popular activation function has a scale equivariant derivative. Yang (2019) shows how to write Batchnorm and Average Pooling as a Tensor Program.

Convolutions. Convolutional layers can be seen as a collection of dense weight matrices where width corresponds to the number of channels (Yang, 2019). With kernel positions ker , input channels [ n l ] and output channels [ n l +1 ] , the weights of a stride-1 convolution are given by { W l iαβ } i ∈ ker,α ∈ [ n l +1 ] ,β ∈ [ n l ] , so that for each i ∈ ker , W l i ∈ R n l +1 × n l is a dense matrix. With { x l iα } i ∈ pos l ,α ∈ [ n l ] , the convolution operation is given by

which performs MatMul and Avg and where ker, pos l 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 ∥ v t ∥ .

Residual connections. A residual connection propagates the current activation forward, skipping an arbitrarily complex nonlinear block f l t : R n l → R n l +1 in between, where f l t can depend on time-dependent parameters like a weight matrix. The forward pass can be written as

If x l t = Θ(1) for all layers l holds in the model without residual connections, it also holds in the model with residual connections. At fixed depth, f l t = o (1) should be avoided, as it would hold that x l +1 t = x l t in the infinite-width limit and the layer would be superfluous. The derivative of the activations becomes

where the second term stays the same as without the residual connection. For the example of f l t being a fully connected layer we get dx l -1 t = dx l t +( W l t ) T ( dx l t ⊙ ϕ ′ ( W l t x l -1 t ) ) . In this example, the derivative with respect to the weights becomes

where the residual connection does not alter the functional dependence on dx l t and x l t compared to a MLP, but implicitly influences the weight gradient since dx l t and x l t are altered. As for the forward pass, the gradient scaling dx l t gets stabilized in the backward pass so that ∂f l t ∂x l -1 t is now allowed to be vanishing with width. Again, we are not aware of an architecture in which that would be desirable. Since a residual connection does not introduce learnable parameters, it interferes in ∥ v t ∥ only implicitly through the stabilized gradients in earlier layers, which can contribute non-vanishingly to ∥ v t ∥ even if later layers are wrongly scaled and their scaling is not adapted.

Adam as a 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 Θ( n -1 ) for hidden and output

layers (Yang et al., 2022). Yang et al. (2023b) provide proofs for arbitrary optimizers that perform generalized, nonlinear outer products. In the example of Adam, the update rule can be written as

where γ i , ω i are the weights that stem from the moving averages. By using a learning rate of n -1 and using the fact that both u and v have approximately iid coordinates of order Θ(1) , the law of large numbers yields Θ(1) updates of the form

Any other learning rate scaling would either result in blowup or vanishing updates.

Adaptive optimizers have not been used for the ascent/perturbation step. In the descent/update step, nothing changes compared to unperturbed optimization as long as we ensure stable perturbations.

F.6 Influence of width-dependent weight multipliers on bcd -parameterizations

Our definition of bcd -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 is that layerwise weight multipliers scale the layerwise gradients. Here, we study how the introduction of weight multipliers affects bcd -parameterizations.

In this section, we consider L -hidden layer MLPs with weight multipliers { a l } l ∈ [ L +1] , width n ∈ N , inputs ξ ∈ R d in , and with outputs f ( ξ ) := n -a L +1 W L +1 x L ( ξ ) where the activations x L ( ξ ) are defined via the iteration

We define abcd -parameterizations in the same way as bcd -parameterizations, but instead of MLPs we use MLPs with weight multipliers { a l } l ∈ [ L +1] .

Definition F.6 ( abcd -parametrization ) . An abcd -parametrization { a l } l ∈ [ L +1] ∪ { b l } l ∈ [ L +1] ∪ { c l } l ∈ [ L +1] ∪ { d l } l ∈ [ L +1] ∪ { d } defines the training of an MLP with weight multipliers { a l } l ∈ [ L +1] with SAM in the following way:

(a) Initialize weights iid as W l ij ∼ N (0 , n -2 b l ) .

(b) Train the weights using the SAM update rule with layerwise learning rates,

with the scaled perturbation ε t via layerwise perturbation radii,

W.l.o.g. we set ∥ v t ∥ = Θ(1) , which prevents nontrivial width-dependence from the denominator. This imposes the constraints:

with at least one equality required to hold, where l ∈ [2 , L ] , and where ∇ x L f = n -a L +1 W L +1 = Θ( n -c ∇ ) with c ∇ = min( b L +1 + a L +1 , c L +1 +2 a L +1 ) . The normalization v t / ∥ v t ∥ removes one degree of freedom from { d l } l ∈ [ L +1] via the equivalence { d ′ l } l ∈ [ L +1] ∼ = { d l } l ∈ [ L +1] iff there exists a C ∈ R such that d ′ l = d l + C for all l ∈ [ L +1] . ◀

F.6.1 abcd -equivalence classes

Update scalings behave as in SGD. The weight multiplier n -a l scales the gradient ∇ W l f by n -a l . In the following forward pass, another multiplication of the weight updates with n -a l leads to the activation update scaling n -2 a l . This can be counteracted by adapting the learning rate scaling. For

abc -parameterizations and SGD training, this induces the layerwise equivalence between parameterizations with ( a l , b l , c l ) or with ( a l + θ l , b l -θ l , c l -2 θ l ) . The extension of all of our results to Adam as a base optimizer is straightforward, since learning rate scalings and perturbation scalings are decoupled. For Adam, c l should be adapted to c l -θ l .

Again, perturbations with joint gradient normalization complicate matters compared to SGD and Adam. Keeping the gradient norm scalings invariant under a l ↦→ a l + θ would require d l ↦→ d l -θ , but keeping the activation perturbation scaling invariant would require d l ↦→ d l -2 θ as for updates. Consequently, an exact equivalence between abcd -parameterizations at finite width requires θ to be the same for all layers and the conflicting gradient norm in the denominator and perturbation scaling in the numerator to be accounted for by adapting the global perturbation scaling d ↦→ d -θ (together with d l ↦→ d l -θ ). In other words, (SAM) with layer-coupling gradient normalization (LP) does not have layerwise analytical equivalence classes at finite width. Below, we provide two 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 proofs are provided at the end of this section.

Lemma F.7 ( abcd -equivalence classes ) . Let f t ( ξ ) denote the output of a MLP in a stable abcd -parameterization with weight multipliers { a l } l ∈ [ L +1] 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 θ ∈ R and any C ∈ R , f t ( ξ ) stays fixed for all t and ξ if, for all l ∈ [ L +1] ,

Remark F.8 ( Infinite-width equivalences ) . In the infinite-width limit, abcd -parameterizations remain equivalent under ( a l + θ l , b l -θ l , c l -2 θ l , d l -2 θ l , d ) layerwise as long as the set of layers that contribute to the gradient norm non-vanishingly remains invariant. The gradient norm constraints for ∥ v l ∥ = O (1) become

where ∇ x L f = n -a L +1 W L +1 = Θ( n -c ∇ ) with c ∇ = min( b L +1 + a L +1 , c L +1 +2 a L +1 ) remains invariant under equivalence transformations. ◀

Remark F.9 ( SAMwith layerwise gradient normalization ) . As the layer coupling is induced by the joint gradient normalization in the perturbations, layerwise gradient normalization simplifies the analysis. For (SAM) with layerwise gradient normalization (LN) of the perturbations global perturbation scaling d is superfluous, and there exist layerwise equivalence classes: For any { θ l } l ∈ [ L +1] ⊂ R ,

To understand this equivalence, observe that any layerwise gradient scaling is cancelled out by the normalization ∇ W l L / ∥∇ W l L∥ . Only the n -a l factor from subsequent forward passes has to be counteracted. ◀

Remark F.10 ( SAM with decoupled perturbation numerator and denominator scaling ) . A perturbation rule with 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 ρ · n -d l and separate gradient norm scaling n -˜ d l , global perturbation scaling d is superfluous, and there exist layerwise equivalence classes: For any { θ l } l ∈ [ L +1] ⊂ R ,

This perturbation rule also allows us to recover an analytical equivalence between trivial weight multipliers a l = 0 for all l , and any other weight multipliers. ◀

F.6.2 µP 2 under non-trivial weight multipliers.

Our goal here is 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 abcd -parameterizations from (LP) is not an issue if we are interested in effective perturbation properties and recovering µ P 2 for arbitrary weight multipliers { a l } l ∈ [ L +1] , as the equivalence breaks due to varying gradient norm contributions, which are inconsequential for achieving effective perturbations.

As we aim to reproduce µ P 2 , we restrict ourselves to the µ 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 ρn -d to at least achieve non-vanishing perturbations in some layers. The following lemma shows even more: The choice

achieves effective perturbations in all layers with the naive (SAM) update rule with naive perturbation scaling ρ · n 0 , 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 conflict 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 1, 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 abcd -parameterization where { ( a l , b l , c l ) } l ∈ [ L +1] are chosen from the µ P equivalence class, and where there is some C ∈ R such that d l = C for all l ∈ [ L + 1] . This reduces to training a MLP with weight multipliers with (SAM) with global perturbation scaling ρn -d for some d ∈ R . Effective perturbations in all layers are achieved and all layers contribute non-vanishingly to the gradient norm if and only if

Achieving µ P 2 with the current implementation of the 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 2 when the goal is simple perturbation scaling.

Lemma F.12 ( Effective perturbations with the mup -package ) . Consider an abcd -parameterization where { ( a l , b l , c l ) } l ∈ [ L +1] are chosen from the µ P equivalence class, and with the weight multipliers a L +1 = I ( l = L +1) as in the mup -package.

(a) ( mup -package global scaling effectively perturbs hidden layers) Under global scaling d l = C , C ∈ R , for all l ∈ [ L +1] , maximal stable perturbations are achieved with d = 0 . In this parameterization, hidden layers are effectively perturbed, but input and output layers are not effectively perturbed.

(b) ( µ P 2 with the mup -package) Effective perturbations in all layers are achieved with the choice d = d 1 = d L +1 = -1 / 2 and d l = 1 / 2 for l ∈ [2 , L ] .

The following lemma covers the general case how to achieve µ P 2 given arbitrary weight multipliers. Lemma F.13 ( µ P 2 with arbitrary weight multipliers ) . Consider an abcd -parameterization where { ( a l , b l , c l ) } l ∈ [ L +1] are chosen from the µ P equivalence class. Then effective perturbations in all layers are achieved with the choice d = min l ∈ [ L +1] ( -a l -1 / 2 I ( l = 1) + 1 / 2 I ( l = L +1)) , and

The following lemma shows that weight multipliers that achieve µ P 2 with naive perturbation scaling under perturbations with layerwise normalization (LN) are exactly the same as the ones for (LP).

Lemma F.14 ((LN) with naive perturbation scaling can effectively perturb all layers ) . Consider (SAM) with layerwise normalization (LN) . Assume { ( a l , b l , c l ) } l ∈ [ L +1] are chosen from the µ P equivalence class, and assume there is some C ∈ R such that d l = C for all l ∈ [ L +1] . Then all layers are effectively perturbed if the multipliers are chosen as

Proof of Lemma F.14. As derived in Appendix F.7, under a l = 0 for all l ∈ [ L +1] , all layers are effectively perturbed if and only if d l = -1 / 2 · I ( l = 1) + 1 / 2 · I ( l = L +1) . Now we can exploit the layerwise equivalence relation to enforce d l = C in each layer by adapting all a l .

Proof of Lemma F.11. In general, in the abc -equivalence class of µ P, the l -th layer's gradient norm is scaled by n -a l . This induces the generalized gradient norm constraints for ∥∇ W L ∥ = Θ(1) ,

Effective perturbations are achieved when ρn -d -d l -a l ∇ W l L = Θ( n -I ( l> 1) ) , which induces the perturbation stability constraints

with effective perturbations whenever the equality of the respective layer holds.

Under global scaling, the gradient norm constraints become, for some C ∈ R ,

and the conditions for effective perturbations become

As d + C is a common term in all layers, we get the relations a l = a 1 +1 / 2 , a L +1 = a 1 +1 , so that all gradient norm constraints are simultaneously satisfied with C = -a l and effective perturbations are achieved in all layers with d = -a l .

Proof of Lemma F.12. Under the choice a l = I ( l = L +1) , the gradient norm constraints become

and the conditions for effective perturbations become

Proof of (a):

Satisfying the gradient norm constraints with global scaling requires d l = 0 for all l ∈ [ L +1] , then the minimal stable choice of d is d = 0 which only effectively perturbs hidden layers.

Proof of (b):

The choice d = -1 / 2 and d 1 = -1 / 2 saturates the gradient norm constraint and achieves effective perturbations in the input layer. Then the choice d l = 1 / 2 and d L +1 = -1 / 2 satisfies the gradient norm constraints and achieves effective perturbations in all layers.

Proof of Lemma F.7. To understand the influence of weight multipliers on updates and perturbations, first note that under an equivalence transformation of all abcd -parameters w.l.o.g from a l = 0 for all l ∈ [ L +1] , the scalings of h l , x l and of n -a l W l remain invariant. This implies that the scalings of ∇ x L f = n -a L +1 W L +1 , ∇ h l f = ∇ x l f ⊙ ϕ ′ ( h l ) and ∇ x l f for all l ∈ [ L ] also remain invariant. Hence the weight gradients, ∇ W L +1 f = n -a L +1 x L and ∇ W l f = n -a l ∇ h l f · ( x l -1 ) ⊤ are scaled by n -a l in each layer.

In the following forward pass, we get

so that activation/output updates and perturbations of layer l are scaled by n -2 a l .

Again, a complication compared to SGD or Adam arises through the gradient normalization of SAM's weight perturbation. If the gradients are simply normalized layerwise ε l = ρ · n -d l ·∇ W l L / ∥∇ W l L∥ , the n -a l -term from the backward pass cancels out, and only in the forward pass we get a scaling n -a l . Hence an exact layerwise equivalence still exists for SAM with layerwise gradient normalization:

Under joint gradient normalization (SAM), as we consider in our definition of bcd -parameterizations, keeping the gradient norm scalings invariant under a l ↦→ a l + θ would require d l ↦→ d l -θ , but keeping the perturbation scaling invariant would require d l ↦→ d l -2 θ as for updates. Consequently, due to the layer coupling of joint gradient normalization ∥∇ W L∥ , an exact equivalence between abcd -parameterizations at finite width requires θ to be the same for all layers and the conflicting gradient norm in the denominator and perturbation scaling in the numerator to be accounted for by d l ↦→ d l -θ and d ↦→ d -θ .

In the infinite-width limit, abcd -parameterizations remain equivalent under ( a l + θ l , b l -θ l , c l -2 θ l , d l -2 θ l , d ) layerwise as long as the set of layers that contributes to the gradient norm nonvanishingly remains invariant. The gradient norm constraints for ∥ v l ∥ = O (1) become

where ∇ x L f = n -a L +1 W L +1 = Θ( n -c ∇ ) with c ∇ = min( b L +1 + a L +1 , c L +1 +2 a L +1 ) remains invariant under equivalence transformations.

Proof of Lemma F.13. First, the choices of d l 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 ∥∇ W L ∥ = Θ(1) -constraints. For input layers, since -d ≥ a 1 + 1 / 2 , it holds that d 1 + a 1 ≥ -1 / 2 . For hidden layers, since -d ≥ a l , it holds that d l + a l ≥ 0 . For output layers, since -d ≥ a L +1 -1 / 2 , it holds that d L +1 + a L +1 ≥ 1 / 2 . Observe that the minimizer in the definition of d saturates its gradient norm constraint.

F.7 The spectral perspective on µP 2

While Tensor Programs allow to track the transformations of vectors like activations, Yang et 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 on the activations, under certain non-cancellation assumptions and limited batch size. For all MLP layers, they show that µ P is equivalent to achieving the condition

at all times t , where W l t : R fan_in → R fan_out . This condition is achieved with initialization σ l , SGD learning rate η l and Adam learning rate η Adam l chosen as,

This generalizes µ P to varying widths inside the network. For varying widths, we adopt the notation W l : R n l -1 → R n l with n 0 = d in and n L +1 = d out , whereas fan_in and fan_out always adapt to the weight matrix under consideration.

To understand why the spectral norm is desirable, note that ∆ W l = η l ∇ h l L ( x l -1 ) ⊤ is low rank and aligned with the incoming activations. For batch size 1 , we even have rank-1 updates with ∥ ∆ W l ∥ ∗ = η l ∥∇ h l L∥ 2 ∥ x l -1 ∥ 2 , aligned with the incoming activations x l -1 , hence ∥ ∆ W l x l -1 ∥ 2 = ∥ ∆ W l ∥ ∗ ∥ x l -1 ∥ 2 . This allows to achieve ∥ ∆ x l ∥ 2 = Θ( √ n l ) irrespective of the layer type with ∥ ∆ W l t ∥ ∗ = Θ( √ n l /n l -1 ) .

Our simple condition that perturbations should scale like updates, which is rigorously justified by our Tensor Program based proof in Appendix E, now allows to derive the correct perturbation scalings using the spectral weight perspective.

Layerwise perturbations. As a simple starting point, consider a variant of (SAM) that does not globally normalize the gradient of all layers jointly, but uses layerwise normalization (LN), resulting in the layerwise perturbation rule,

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 are not coupled, and the spectral condition ∥ ε l ∥ ∗ = Θ( √ fan_out / fan_in ) immediately requires choosing

Perturbations with global gradient normalization. Perturbations that are globally normalized across layers have usually been implemented practice according to the GitHub repositories provided for effective perturbations in layer l with width-independent hyperparameter ρ ≥ 0 .

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 it 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 widthindependence 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 to scale each denominator term separately from the corresponding numerator term (DP), the perturbation radius in each layer for the numerator can be scaled like updates, ρ l = Θ ( fan_out fan_in ) .

Now, to ensure Θ(1) in the denominator, each input and hidden-like gradient norm ∥∇ W l L∥ F , l ∈ [ L ] , achieves width-independence if it scaled by √ fan_out / fan_in . The same rule applies to biases when understanding them as weights R → R n l to the input 1 . These scalings are derived in the next paragraph. The last-layer gradient norm ∥∇ W L +1 L∥ F should be scaled as ( n L n L +1 ) -1 / 2 , and ∥∇ b L +1 L∥ F as n -1 / 2 L +1 .

If we care about the correct width-independent constants, observe that the learning rate scaling η L +1 = Θ( fan_out / fan_in ) induces ∥ W L +1 ∥ = ∥ ∆ W L +1 ∥ = Θ( √ n 3 L +1 /n L ) . If we wanted to achieve ∆ W L +1 = Θ( √ n L +1 /n L ) we would need η L +1 = Θ(1 / fan_in ) . As n L +1 = d out is width-independent, √ fan_out / fan_in would result in the same width-dependent scaling for the last layer, but ignoring large constants can introduce a significant width-independent spectral distortion. For example in ImageNet1K, n L +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.

Derivation of gradient norm ∥∇ W l L∥ F scalings. In this paragraph, ∥ · ∥ may denote the Frobenius or spectral norm. As all matrices are of limited rank, both norms scale equivalently. As a first step, ∥∇ h l L∥ = Θ( 1 √ n l ) can be reconstructed from

Now, for input and hidden layers, ∥∇ W l L∥ = ∥∇ h l L∥∥ x l -1 ∥ 2 = Θ( √ n l -1 /n l ) . Multiplying by the inverse yields width-independent scaling. The output layer gradient ∇ W L +1 L ∈ R n L +1 × n L is given by ( ∇ W L +1 L ) ij = x L j = Θ(1) , so that ∥∇ W L +1 L∥ = Θ( √ n L n L +1 ) . Biases before the last layer follow the scheme ∥∇ b l L∥ = ∥∇ h l L∥ = Θ( √ 1 /n l ) = Θ( √ fan_in / fan_out ) . The last layer bias ∥∇ b L +1 L∥ = √ n L +1 scales width-independently as it should, but needs to be scaled by a different constant 1 / √ fan_out than earlier layers.

Extensions to ASAM. As ASAM cannot be written as a NE ⊗ OR ⊤ program, its scaling can only be derived heuristically. As provided in Table 1 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 √ n L , 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.

Due to its small initialization, it holds that ∥ W L +1 ∥ 2 F = ∥ ∆ W L +1 ∥ 2 F = Θ( n 3 L +1 n L ) . For perturbations that fulfill the spectral condition ρ L +1 ∥ W L +1 ∥ 2 F ∥∇ W l +1 L∥ ∗ = Θ( √ n L +1 n L ) , we need to choose ρ L +1 = ρ · 1 n 3 L +1 (width-independent, but very small). The last-layer denominator term scales as ∥ W L +1 ∥ F ∥∇ W l +1 L∥ ∗ = Θ( √ n 3 L +1 n L · √ n L n L +1 ) = Θ( n 2 L +1 ) , which is width independent, but

For layerwise ASAM, first note that ∥ W l t ∥ F = Θ( ∥ W l 0 ∥ F ) = Θ( √ n l ) for input and hidden layers l ∈ [ L ] . As the numerator contains ∥ W l t ∥ 2 F , it requires the layerwise perturbation scaling 1 fan_in . In the denominator, width independence is achieved with the multiplier √ 1 fan_in , since ∥ W l ∥ F ∥∇ W l L∥ ∗ = √ n l √ n l -1 n l = √ n l -1 . Again, the output layer requires a special treatment.

can be a large constant, as for ImageNet1K. The output bias numerator exactly conforms with the correct scaling ∥∇ b L +1 L∥ 2 F = n L +1 = fan_out / fan_in = Θ(1) .

Note that weight decay may break statements like ∥ W l t ∥ F = Θ( ∥ W l 0 ∥ F ) 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 -package introduces the output layer weight multiplier n -1 L 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 -1 L W L +1 and n -1 L ∆ W L +1 remain the unique ones that achieve µ P, but ∇ W L +1 L is scaled by n -1 L . This requires adapting the last-layer learning rate η L +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 ρ L +1 = ρ 1 = ρ · n and ρ l = ρ for l ∈ [2 , L ] , and the gradient norm contributions should be scaled by √ n for input and output layers, and by 1 for hidden layers. The Tensor Program perspective on weight multipliers can be found in Appendix F.6. The correct width-independent constants are achieved with the last-layer numerator scaling ρ L +1 = ρ · n L and the last-layer denominator scaling √ n L /n L +1 , since ∇ W L +1 L = Θ( √ n L +1 /n L ) and for the numerator we get an additional n -1 L in the forward pass.

For SAM-ON nothing changes, as only input-like layers are perturbed. For elementwise ASAM, ignoring width-independent constants, nothing changes as the weight multiplier n -1 L increases the weight scaling W L +1 and decreases the gradient scaling ∇ W L +1 L by the same amount. The additional n -1 L -factor in the numerator is cancelled out by the additional W L +1 -factor. For the correct width-independent constants with decoupled numerator and denominator scaling, we would scale the denominator by √ n L /n 3 L +1 with or without weight multiplier, and scale the numerator by ρ L +1 = ρ · n L /n 2 L +1 with or without weight multiplier. For the example of layerwise ASAM, we still get for the denominator ∥ W L +1 ∥ F ∥∇ W l +1 L∥ ∗ = Θ( n 2 L +1 ) , again because the weights W L +1 are scaled up by n L 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 n -1 L in the subsequent forward pass, leading to an unchanged ρ L +1 = ρ · n -3 L +1 , which is width-independent but potentially leads to numerical issues.

Code for µ P 2 with separate denominator scalings. Algorithm 1 provides a PyTorch code example that implements the above µ P 2 scalings for SAM, scaling the gradient norm contributions of all layers to Θ(1) (equivalent to ( a -µP 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 × constant or constant × n and include input and output weights. The last-layer multiplier n -1 L is chosen so that input and output layers can be scaled by the same width-dependent factor. On the other hand, 'matrix-like ´ parameters scale as n × n and include hidden weights. The implementation uses a base width at which µ P 2 and SP are equivalent; all width-dependent scalings then scale with width-multipliers width / base_width . This allows to immediately transfer well-performing settings from SP to µ P 2 .

Let us recapitulate how the µ P 2 scaling in the following code arises. The crucial variables to track are factor , group["rho"] and group["gradnorm_scaling"] . For limited batch size, the spectral and Frobenius norm of gradients scale equivalently, and we get, for all l ∈ [ L ] ,

We want to scale each weight's contribution in the denominator to be width-independent, hence need the factor √ factor with factor = fan_out / fan_in . For the numerator, the spectral condition ( ∗ ) demands ∥ ρ l · ∇ W l L∥ ∗ ! = Θ( √ fan_out fan_in ) , so that we need to scale the weight's perturbation radius to ρ l = ρ · factor . Since the 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).

Algorithm 1: Pytorch implementation of µ P 2 for SAM using the mup -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["gradnorm_scaling"] in the denominator to be width-independent. The numerator terms group["rho"] of all weight tensors are scaled to achieve effective perturbations. This scaling is equivalent to ( a -µP 2 ) 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 bcd -parameterizations to 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 a base width at which SP and µ P are equivalent, allowing to immediately transfer setups that perform well in SP. We use the mup -package only for ViTs, and our implementation of µ P 2 resembles the pseudocode provided in Algorithm 1. For ResNets, we use no width-dependent lastlayer multiplier. At initialization, µ P differs from SP only through a smaller last layer initialization. For MLPs we exactly implement the bcd -parameterization with b L +1 = 1 , but use the large widthindependent input layer initialization variance 2 instead of the width-independent 2 /d in in µ P, which can be seen as a tuned initialization variance multiplier. For ResNets and Vits, we initialize the last layer to 0 in µ P, which corresponds to b L +1 →∞ and which recovers the limit behaviour f 0 → 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 to mimic standard training procedures. For µ P, we resort to the optimal multipliers from Yang et al. (2022). We then find the optimal learning rate and perturbation radius for each bcd -parametrization and SAM variant separately.

ResNets. For ResNet18 experiments, we augment the CIFAR10 data with random crops and random horizontal flips, set labelsmoothing to 0 . 1 and use a cosine learning rate schedule. ResNets in µ P have base width 0 . 5 , gradient norm scaling according to Definition 4 and their last layer is initialized to 0 . 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 µ 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 , 2 -7 , . . . , 2 8 ] , weight decay [0 , 10 -5 , 10 -4 , 5 · 10 -4 , 10 -3 , 10 -2 ] and momentum [0 , 0 . 1 , 0 . 4 , 0 . 7 , 0 . 9] . Learning rate and perturbation radius grids were either set to [2 -10 , 2 -9 , . . . , 2 1 ] or centered around recommendations from the literature. The optimal hyperparameter configurations found from at least 150 runs for each SAM variant are summarized in Table G.1. Learning rates and perturbation radii were further tuned with the experiments from Appendix H.3.3.

ViTs. We train ViT-S/16 with 6 layers and 12 attention heads on ImageNet1K (Deng et al., 2009) and a ViT-S/4 with 12 layers and 12 attention heads on CIFAR100 (Krizhevsky et al., 2009) (see Appendix H.6), again adopting the hyperparameter settings from Müller et al. (2024). This means we use AdamW as a base optimizer with warmup and a cosine learning rate decay. For CIFAR100, we use random crops, random horizontal flips and AutoAugment as data augmentations. For Imagenet we use the original preprocessing from Huggingface vit-base-patch16-224 (Wu et al., 2020). For µ P, we tune multipliers at a basewidth 384 , initialize the last layer and query weights to 0 . By

using the µ P package, the relative perturbation scalings change as 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 CIFAR100 and ImageNet are summarized in Table G.2. For µ P, the learning rate, perturbation radius, input multiplier, output multiplier and weight decay were tuned using 3 independent runs of Nevergrad NGOpt with budget 56 on ImageNet. The same multipliers are used on CIFAR100.

Figures. Whenever multiple runs with independent random seeds are used for training, confidence bands cover the interval from the empirical 2 . 5% - to the empirical 97 . 5% -quantile. The line then denotes the average of all runs. When confidence bands are given, but the number of independent runs is not specified, the number of 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 a single GPU, our µ P 2 -SAM training script for MLPs of width 4096 on CIFAR10 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 in Table G.3.

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 = d l = 0 for all l ). Global perturbation scaling (global) denotes the maximal stable scaling n -d of the global perturbation radius that achieves effective perturbations in some layers without layerwise perturbation scaling ( d l = 0 for all l ).

H.1 SAM is approximately LL-SAM in µ P with global perturbation scaling

Figure H.1 compares SAM in µ P under global perturbation scaling ( µ P-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 is closely related to 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 1 / 2 , 1 and n -1 / 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 µ P-global SAM than SGD on all of the tracked statistics, in particular at large widths. Only perturbing the last layer still affects the gradients in 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 µ P 2 does improve over SGD for all widths (Figure H.3). Last-layer perturbation norms coincide by design with the global perturbation radius n -d ρ 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

where the loss derivative χ t always cancels out due to the normalization and the global gradient norm ∥ v t ∥ is dominated by the last-layer gradient norm due to the global scaling (Theorem 11). 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 are quite noisy. The last-layer Hessian spectral norm explodes with width in µ P, because last-layer learning rate is scaled as n -1 , hence the edge of stability explodes. ResNets in µ P are more stable, their Hessian spectral norm even shrinks with width (not shown).

Contrast the results for µ P-global with the results for µ P 2 in Figure H.2 for a comparison with SGD in µ P. The Hessian spectral norm is reduced by SAM as you would expect. Additionally µ P 2 shows low variability in performance and all other statistics. SAM in µ P 2 does not reduce the feature rank compared to SGD in µ P. This suggests that the conclusions drawn by Andriushchenko et al. (2023a) do not apply to MLPs in µ P.

H.2 Propagating perturbations from the first layer does not inherit SAM's benefits

Here we apply a parametrization that only effectively perturbs the first layer weights (derived in Example F.1). Figure H.2 shows that effective first-layer SAM loses both µ P 2 SAM's improvement in test accuracy as well as SAM's inductive bias towards smaller gradient norm and Hessian norm, i.e. lower sharpness in MLPs. This performance deterioration occurs although the perturbation of first-layer SAM has an effect of the same order of magnitude as µ P 2 on weight and activation updates in all layers. This shows that mere propagation of weight perturbations from earlier layers cannot replace effective weight perturbations in each layer in order to 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 µ P 2 , 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.

64

256

1024

4096

16384

0

.

00000

0

.

00025

0

.

00050

0

.

00075

0

.

00100

0

.

00125

0

.

00150

0

.

00175

‖

˜

W

-

W

‖

F

Input layer

64

256

1024

4096

16384

0

.

000

0

.

002

0

.

004

0

.

006

0

.

008

Hidden layer

64

256

1024

4096

16384

0

.

000

0

.

005

0

.

010

0

.

015

0

.

020

0

.

025

Output layer

SAM-SGD

64

256

1024

4096

16384

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

1

.

4

‖

˜

W

-

W

‖

F

/

‖

W

‖

∗

×

10

-

5

64

256

1024

4096

16384

0

.

0000

0

.

0002

0

.

0004

0

.

0006

0

.

0008

0

.

0010

0

.

0012

0

.

0014

64

256

1024

4096

16384

0

.

000

0

.

005

0

.

010

0

.

015

0

.

020

0

.

025

SAM-SGD

64

256

1024

4096

16384

9

.

25

9

.

50

9

.

75

10

.

00

10

.

25

10

.

50

10

.

75

||

∆

W

l

||

∗

/c

l

Input layer

64

256

1024

4096

16384

6

7

8

9

10

Hidden layer

64

256

1024

4096

16384

6

8

10

12

14

Output layer

SAM-SGD

SGD

64

256

1024

4096

16384

0

.

12

0

.

14

0

.

16

0

.

18

0

.

20

Coord.wise

‖

x

t

-

x

0

‖

Input layer

64

256

1024

4096

16384

0

.

15

0

.

20

0

.

25

0

.

30

Hidden layer

64

256

1024

4096

16384

3

4

5

6

7

8

Output layer

SAM-SGD

SGD

64

256

1024

4096

16384

100

200

300

400

500

Feature rank

Input layer

64

256

1024

4096

16384

100

200

300

400

500

Hidden layer

64

256

1024

4096

16384

8

.

0

8

.

2

8

.

4

8

.

6

8

.

8

9

.

0

Output layer

SAM-SGD

SGD

64

256

1024

4096

16384

0

2

4

6

8

10

12

14

16

Gradient norm

SAM-SGD

SGD

64

256

1024

4096

16384

0

500

1000

1500

2000

2500

3000

3500

Hessian spectral norm

SAM-SGD

SGD

64

256

1024

4096

16384

0

10000

20000

30000

40000

Hessian trace

SAM-SGD

SGD

64

256

1024

4096

16384

60

65

70

75

80

85

90

95

100

Training accuracy

SAM-SGD

SGD

64

256

1024

4096

16384

46

48

50

52

54

56

58

60

Test accuracy

SAM-SGD

SGD

Gradient norm

64

256

1024

4096

0

2

4

6

8

10

12

14

‖

˜

W

-

W

‖

F

Input layer

64

256

1024

4096

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

Hidden layer

64

256

1024

4096

0

.

0000

0

.

0025

0

.

0050

0

.

0075

0

.

0100

0

.

0125

0

.

0150

0

.

0175

Output layer

perturbation

mpp

first layer

64

256

1024

4096

0

.

0000

0

.

0025

0

.

0050

0

.

0075

0

.

0100

0

.

0125

0

.

0150

0

.

0175

‖

˜

W

-

W

‖

F

/

‖

W

‖

∗

64

256

1024

4096

0

.

000

0

.

005

0

.

010

0

.

015

0

.

020

64

256

1024

4096

0

.

000

0

.

005

0

.

010

0

.

015

0

.

020

0

.

025

perturbation

mpp

first layer

64

256

1024

4096

11

.

5

12

.

0

12

.

5

13

.

0

||

∆

W

l

||

∗

/c

l

Input layer

64

256

1024

4096

6

8

10

12

Hidden layer

64

256

1024

4096

4

6

8

10

12

Output layer

mpp

first layer

SGD

64

256

1024

4096

0

.

14

0

.

16

0

.

18

0

.

20

0

.

22

0

.

24

0

.

26

Coord.wise

‖

x

t

-

x

0

‖

Input layer

64

256

1024

4096

0

.

15

0

.

20

0

.

25

0

.

30

0

.

35

Hidden layer

64

256

1024

4096

2

3

4

5

6

7

8

9

Output layer

mpp

first layer

SGD

64

256

1024

4096

100

200

300

400

500

Feature rank

Input layer

64

256

1024

4096

100

200

300

400

500

Hidden layer

64

256

1024

4096

8

.

0

8

.

2

8

.

4

8

.

6

8

.

8

9

.

0

Output layer

mpp

first layer

SGD

64

256

1024

4096

10

20

30

40

50

mpp

first layer

SGD

64

256

1024

4096

0

200

400

600

800

1000

1200

Hessian spectral norm

mpp

first layer

SGD

64

256

1024

4096

0

250

500

750

1000

1250

1500

1750

Hessian eigenvalue gap

mpp

first layer

SGD

64

256

1024

4096

60

70

80

90

100

Training accuracy

mpp

first layer

SGD

64

256

1024

4096

44

46

48

50

52

54

56

58

Validation accuracy

mpp

first layer

SGD

H.3 Hyperparameter transfer

In this section, we provide supplemental evidence that, µ P 2 is the unique perturbation scaling that robustly achieves hyperparameter transfer in µ 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 in µ P

Figure H.3 shows that in µ P 2 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 µ P 2 outperforms SAM in µ P with global perturbation scaling at 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 is 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 ≥ 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

100

90

90

Training accuracy

Training accuracy

80

70

70

60

60

50

50

40

10-1

100

10-1

100

Learning rate

Perturbation radius

60.0

60

width

57.5

58

64

Validation accuracy

55.0

Validation accuracy

64

256

1024

52.5

4096

50.0

47.5

45.0

48

42.5

46

40.0

10-1

100

10-1

100

Learning rate

Perturbation radius

100

100

95

95

90

90

Training accuracy

Training accuracy

85

85

80

75

70

65

65

60

60

55

55

10-1

100

10-1

100

Learning rate

Perturbation radius

58

58

width

64

Validation accuracy

Validation accuracy

64

256

1024

4096

2

50

50

48

48

10-1

100

10-1

100

Learning rate

Perturbation radius

100

100

95

95

90

06

Training accuracy

Training accuracy

85

85

80

80

75

75

70

70

65

65

60

60

55

55

101

102

10-1

100

Learning rate

Perturbation radius

58

58

width

64

Validation accuracy

56

Validation accuracy

256

1024

54

4096

2

52

50

50

48

48

101

102

10-1

100

Perturbation radius

Learning rate

Figure H.4 shows that global perturbation scaling does transfer the same perturbation instability threshold, whereas in µ P-naive every fixed perturbation radius becomes unstable at sufficient width (Figure H.7). But in µ P-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 µ P-global but with a 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.

10

-

3

10

-

2

10

-

1

Perturbation radius

ρ

10

-

1

10 0

Learning rate

η

µ

P-naive

0.0

0.001

0.0018

0.0032

0.0056

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

Perturbation radius

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

Learning rate

43

43

43

43

43

43

43

43

42

41

38

10

46

46

46

46

46

46

46

46

46

45

42

10

50

50

50

50

50

50

50

50

49

49

47

10

54

54

54

54

54

54

54

54

53

53

51

10

58

58

58

58

58

58

58

57

57

56

55

10

60

60

60

60

61

60

61

60

59

59

57

10

61

62

62

61

61

62

62

60

60

58

56

10

59

60

60

60

59

60

59

58

58

55

24

10

54

54

53

56

56

52

54

54

51

49

21

10

width=64

0.0

0.001

0.0018

0.0032

0.0056

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

Perturbation radius

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

Learning rate

49

49

49

49

49

49

49

48

47

44

10

10

54

54

54

55

55

54

54

53

52

50

10

10

61

61

61

61

61

61

61

60

59

57

10

10

68

68

68

68

69

69

69

68

67

65

10

10

74

75

75

76

76

77

76

76

74

70

10

10

80

81

81

82

82

81

81

80

77

73

10

10

80

83

82

81

81

82

82

80

76

68

10

10

79

79

80

81

81

79

76

78

73

47

10

10

72

69

72

72

71

71

69

68

63

24

10

10

width=256

0.0

0.001

0.0018

0.0032

0.0056

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

Perturbation radius

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

Learning rate

53

53

53

53

52

52

51

49

44

11

11

11

60

60

60

60

60

59

58

57

53

11

11

11

69

69

69

69

69

69

68

66

64

11

11

11

77

78

78

78

79

78

77

75

73

11

10

10

83

86

85

85

85

85

85

83

79

10

10

10

90

88

91

91

92

91

88

86

81

10

10

10

95

92

96

94

91

89

92

90

80

10

10

10

87

86

87

86

89

83

86

85

10

10

10

10

84

85

82

81

83

78

73

75

10

10

10

10

width=1024

0.0

0.001

0.0018

0.0032

0.0056

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

Perturbation radius

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

Learning rate

54

54

54

53

53

51

50

15

11

12

12

12

62

62

62

61

61

59

58

12

12

12

12

11

72

73

73

72

72

71

70

12

12

11

11

11

81

82

82

82

82

81

79

12

11

11

11

10

90

88

88

87

86

86

84

11

11

10

10

10

96

95

95

92

92

91

89

11

10

10

10

10

99

98

94

93

94

97

95

10

10

10

10

10

99

94

92

94

92

94

92

10

10

10

10

10

91

85

85

85

87

87

80

10

10

10

10

10

width=4096

52

54

56

58

60

Training accuracy

74

76

78

80

82

Training accuracy

86

88

90

92

94

Training accuracy

90

92

94

96

98

Training accuracy

0.0

0.001

0.0018

0.0032

0.0056

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

Perturbation radius

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

Learning rate

42

42

42

42

42

42

42

42

41

40

37

15

45

45

45

45

45

45

44

44

44

43

41

13

47

47

47

47

47

47

47

47

46

45

44

11

49

49

49

49

49

49

48

48

48

48

46

10

49

49

50

50

50

49

49

49

49

49

47

10

50

50

50

50

49

49

49

49

49

49

48

10

50

50

49

50

49

49

49

49

49

48

47

10

49

48

49

48

48

49

48

48

48

47

22

10

47

47

47

47

46

47

47

47

46

44

20

10

width=64

0.0

0.001

0.0018

0.0032

0.0056

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

Perturbation radius

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

Learning rate

47

47

47

47

47

47

47

46

46

43

19

13

49

50

50

50

50

50

49

49

48

47

17

11

51

51

51

51

51

51

51

51

51

50

13

11

52

52

52

52

52

52

52

52

52

51

11

10

52

52

53

53

53

53

53

53

52

52

10

10

53

53

53

53

53

53

53

53

53

52

10

10

52

52

52

52

52

52

52

52

52

51

10

10

51

52

52

51

51

51

51

51

50

35

10

10

50

50

50

49

49

49

49

49

47

21

10

10

width=256

0.0

0.001

0.0018

0.0032

0.0056

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

Perturbation radius

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

Learning rate

49

49

50

49

49

49

48

47

44

23

14

11

52

52

52

52

52

52

52

51

49

20

12

11

53

53

53

54

54

54

54

53

53

15

11

11

54

54

54

54

55

55

54

54

53

12

11

11

54

54

54

54

55

55

55

55

54

11

11

11

54

54

55

55

55

55

55

55

54

11

11

11

55

55

55

55

55

55

55

55

53

11

11

11

54

54

54

54

54

54

54

53

11

11

11

11

52

52

51

52

51

52

51

49

11

11

11

11

width=1024

0.0

0.001

0.0018

0.0032

0.0056

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

Perturbation radius

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

Learning rate

50

50

50

50

50

49

48

34

28

17

12

11

53

53

53

53

53

53

52

34

24

13

11

10

54

54

54

55

55

55

55

34

17

12

10

10

54

55

55

55

55

55

55

32

13

11

10

10

54

55

55

55

55

55

55

23

11

10

10

10

56

56

56

56

56

56

56

14

10

10

10

10

58

57

57

57

57

57

56

11

10

10

10

10

57

57

56

56

57

56

55

10

10

10

10

10

54

54

54

54

54

53

51

10

10

10

10

10

width=4096

45

46

47

48

49

Validation accuracy

49

50

51

52

53

Validation accuracy

51

52

53

54

55

Validation accuracy

53

54

55

56

57

Validation accuracy

0.0

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

1.7783

3.1623

Perturbation radius

0.005

0.0097

0.0188

0.0364

0.0706

0.1369

0.2657

0.5154

1.0

Learning rate

39

39

39

39

39

39

39

38

37

35

11

10

43

43

43

43

43

43

42

42

41

39

10

10

46

46

47

47

47

46

46

46

46

44

10

10

51

51

51

51

51

51

51

51

50

49

10

10

55

55

55

55

55

56

55

55

55

54

10

10

59

59

59

59

59

59

59

59

59

57

10

10

61

62

61

61

61

60

61

60

60

58

10

10

60

59

60

60

60

59

59

58

58

54

10

10

54

54

54

55

55

54

54

53

51

48

10

10

width=64

0.0

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

1.7783

3.1623

Perturbation radius

0.005

0.0097

0.0188

0.0364

0.0706

0.1369

0.2657

0.5154

1.0

Learning rate

44

44

44

44

44

44

44

44

43

41

33

10

49

49

49

49

49

49

49

48

48

46

41

10

55

55

55

55

55

55

55

54

54

53

50

10

62

62

62

62

62

63

63

62

62

61

59

10

71

71

71

71

72

72

72

72

71

70

67

10

79

79

79

79

79

80

80

79

78

76

69

10

80

82

81

84

80

82

83

82

81

77

49

10

83

81

79

81

78

79

78

79

77

72

10

10

72

71

71

72

72

71

70

70

67

61

10

10

width=256

0.0

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

1.7783

3.1623

Perturbation radius

0.005

0.0097

0.0188

0.0364

0.0706

0.1369

0.2657

0.5154

1.0

Learning rate

47

47

47

47

47

46

46

46

45

43

36

11

52

52

52

52

52

52

52

51

50

49

44

11

61

61

61

61

61

61

60

60

59

57

54

11

71

72

72

72

72

72

72

71

71

69

67

11

81

81

82

82

82

83

82

82

82

80

77

11

88

88

89

89

89

88

90

88

88

85

81

10

96

93

94

95

91

92

92

91

91

91

83

10

89

92

91

93

89

90

88

90

86

87

31

10

83

85

84

84

79

79

80

83

81

75

10

10

width=1024

0.0

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

1.7783

3.1623

Perturbation radius

0.005

0.0097

0.0188

0.0364

0.0706

0.1369

0.2657

0.5154

1.0

Learning rate

47

47

47

47

47

47

47

46

45

43

37

13

54

54

54

54

53

53

53

52

51

50

45

11

63

63

63

63

63

63

62

62

61

59

56

12

75

75

75

75

75

75

75

75

74

73

70

12

85

85

85

85

85

85

85

85

84

82

80

11

92

91

92

93

92

91

92

88

89

87

85

10

97

97

98

97

97

96

94

95

90

91

89

10

98

98

98

94

95

96

92

95

96

90

10

10

81

80

85

92

86

82

88

81

80

78

10

10

width=4096

52

54

56

58

60

Training accuracy

74

76

78

80

82

Training accuracy

86

88

90

92

94

Training accuracy

90

92

94

96

98

Training accuracy

0.0

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

1.7783

3.1623

Perturbation radius

0.005

0.0097

0.0188

0.0364

0.0706

0.1369

0.2657

0.5154

1.0

Learning rate

39

39

39

39

39

38

38

38

37

35

18

14

42

42

42

42

42

42

42

41

40

39

18

13

45

45

45

45

45

45

45

44

44

43

18

11

47

47

47

47

47

47

47

47

46

46

16

10

49

49

49

49

49

49

49

48

48

47

12

10

49

50

49

49

49

49

49

49

49

48

10

10

50

50

50

50

49

50

49

49

49

48

10

10

49

49

49

49

49

49

49

48

47

47

10

10

47

47

47

47

47

47

47

46

45

43

10

10

width=64

0.0

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

1.7783

3.1623

Perturbation radius

0.005

0.0097

0.0188

0.0364

0.0706

0.1369

0.2657

0.5154

1.0

Learning rate

44

44

44

44

44

44

44

43

42

41

33

18

47

47

47

47

47

47

47

46

46

45

41

18

50

50

50

50

50

50

50

50

49

49

47

15

51

51

51

51

52

52

52

52

52

51

50

12

52

52

52

52

52

52

53

52

52

52

51

10

53

53

53

53

53

53

53

53

53

52

51

10

52

53

53

52

53

53

52

53

52

52

37

10

52

52

52

51

51

51

51

51

52

51

10

10

50

50

50

50

50

50

49

50

48

47

10

10

width=256

0.0

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

1.7783

3.1623

Perturbation radius

0.005

0.0097

0.0188

0.0364

0.0706

0.1369

0.2657

0.5154

1.0

Learning rate

46

46

46

46

46

46

46

45

44

43

35

24

49

49

49

49

49

49

49

49

48

47

44

24

52

52

52

52

52

52

52

52

52

51

50

20

53

53

53

53

53

53

54

54

54

54

53

15

54

54

54

54

54

54

54

54

55

54

54

12

54

54

54

54

54

55

55

55

55

55

54

11

55

55

55

55

55

55

55

55

55

55

54

11

55

54

54

54

54

54

54

55

54

54

24

11

51

52

52

52

51

51

52

51

51

50

11

11

width=1024

0.0

0.01

0.0178

0.0316

0.0562

0.1

0.1778

0.3162

0.5623

1.0

1.7783

3.1623

Perturbation radius

0.005

0.0097

0.0188

0.0364

0.0706

0.1369

0.2657

0.5154

1.0

Learning rate

47

47

47

47

47

47

46

46

45

43

36

29

50

50

50

50

50

50

50

50

49

48

45

29

53

53

53

53

53

53

53

53

53

52

51

26

54

54

54

54

54

54

54

54

55

55

54

18

54

54

54

54

55

55

55

55

55

55

55

13

55

55

55

55

55

55

55

56

56

55

55

11

57

57

57

56

57

57

56

57

57

57

56

11

58

58

57

57

57

56

56

56

56

56

11

11

54

54

54

54

54

54

53

54

54

52

11

11

width=4096

46

47

48

49

50

Validation accuracy

49

50

51

52

53

Validation accuracy

51

52

53

54

55

Validation accuracy

54

55

56

57

58

Validation accuracy

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

49

49

49

49

49

47

42

13

52

53

52

52

52

51

45

14

56

56

56

56

56

54

48

13

59

59

59

59

59

57

51

15

61

61

61

61

61

59

53

21

61

61

61

59

60

59

55

25

58

58

58

57

58

59

55

27

55

54

55

54

52

53

51

29

45

44

45

46

44

45

41

19

width=64

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

59

59

59

59

58

55

47

18

66

66

66

65

64

60

51

17

73

73

73

72

71

67

55

17

77

78

78

77

77

74

60

23

81

83

80

79

80

79

65

29

82

79

82

82

82

82

70

33

79

76

80

79

78

80

73

39

69

70

70

72

72

72

70

32

55

52

55

56

56

56

55

23

width=256

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

67

67

67

66

64

59

49

22

76

75

75

74

73

66

53

20

83

83

83

83

81

75

58

18

88

88

88

88

88

85

64

26

91

92

95

90

94

91

73

33

94

94

93

90

91

92

83

38

87

89

90

90

89

93

90

44

76

84

81

78

81

84

89

38

64

65

66

65

66

71

70

27

width=1024

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

71

70

70

69

67

60

49

22

79

79

79

78

76

69

54

19

87

87

86

86

84

78

59

18

93

93

92

91

88

88

66

29

95

94

96

94

91

93

76

34

98

99

96

97

95

97

88

40

91

96

93

93

95

98

96

45

87

82

84

80

89

91

97

42

66

60

72

70

77

79

84

30

width=4096

52

54

56

58

60

Training accuracy

74

76

78

80

82

Training accuracy

86

88

90

92

94

Training accuracy

90

92

94

96

98

Training accuracy

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

46

46

46

46

46

46

42

15

48

48

48

48

48

48

44

15

49

49

49

49

49

49

46

15

49

49

50

49

50

50

48

15

50

50

50

50

50

50

49

22

49

49

49

50

50

50

49

25

49

48

48

48

49

49

49

29

46

46

47

46

46

47

46

30

41

43

42

43

42

41

39

20

width=64

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

51

51

51

51

51

51

47

21

52

52

52

52

53

53

49

20

52

52

53

53

53

54

51

19

52

52

52

53

53

54

53

23

52

52

53

53

53

54

54

29

52

52

52

53

53

54

54

33

51

51

51

51

52

53

54

38

50

49

49

50

50

51

51

33

45

44

45

44

44

46

46

24

width=256

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

53

53

53

53

54

53

48

28

53

53

53

53

54

55

51

27

54

54

54

54

55

56

53

25

54

54

54

54

55

56

55

26

54

54

55

55

55

57

57

33

55

55

55

55

55

56

57

38

54

53

54

54

54

55

57

43

51

52

52

51

52

53

55

37

46

46

46

47

47

48

49

27

width=1024

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

54

54

54

54

54

54

48

31

54

54

54

54

55

56

51

32

54

54

54

55

55

57

53

31

55

55

55

55

55

57

56

29

56

56

56

56

56

58

58

34

58

58

58

57

57

58

59

40

56

56

57

57

56

58

59

44

53

54

54

54

53

56

58

40

47

48

47

47

48

50

53

31

width=4096

46

47

48

49

50

Test accuracy

50

51

52

53

54

Test accuracy

53

54

55

56

57

Test accuracy

54

55

56

57

58

Test accuracy

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-width theory. For SGD, this has originally been observed in GitHub issue 52 of the mup -package 4 . We train MLPs with SAM in variants of SP in Figure H.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 et al. (2022) use weight multipliers tuned at base width 256 and normalize the initialization variance to be invariant to these weight multipliers, according to the Jupyter notebook 5 provided for reproducing their experiments. In addition, they initialize the last layer to 0 which contradicts 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-10. With tuned multipliers but initialization that is invariant to 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 a matter of ongoing work to understand whether this stability of SP is a finite-width or a 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.1), we expect the same shrinkage for SAM in SP in such settings.

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 µ P through a larger output layer initialization and key-query normalization by √ d in SP versus by d in µ P. Here however we even observe transfer in our stricter understanding of SP without any layerwise learning rates or weight multipliers. This is not a peculiarity of SAM; we observe the same learning rate transfer in plain SGD without any momentum or weight decay for MLPs on CIFAR-10 (not shown).

10

-

2

10

-

1

10 0

Perturbation radius

ρ

10

-

2

10

-

1

Learning rate

η

SP-naive, no mults

width

4096

1024

256

10

-

2

10

-

1

10 0

Perturbation radius

ρ

10

-

1

10 0

SP-naive (Yang mults)

10

-

3

10

-

2

10

-

1

Perturbation radius

ρ

10

-

4

10

-

3

10

-

2

10

-

1

SP-naive, initnorm

10

-

3

10

-

2

10

-

1

Perturbation radius

ρ

10

-

4

10

-

3

10

-

2

10

-

1

SP-naive, initnorm, ll=0

ResNets in 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.

(a) No momentum, weight decay or labelsmoothing

(b) Tuned momentum and weight decay

100

100

100

100

Training accuracy

98

Training accuracy

99

Training accuracy

99

Training accuracy

99

96

98

98

98

94

97

97

9

96

96

96

92

102

10 ~1

100

95

10 3

10 2

10 1

95

10 2

101

100

95

10 ~3

10 2

10 1

Learning rate

Perturbation radius

Learning rate

Perturbation radius

Validation accuracy

92

94

93

Validation accuracy

94

93

95

width

0.5

2.0

1.0

Validation accuracy

96.0

9°96

97.0

96.5

Validation accuracy

97.0

96.5

96.0

9°96

width

0.5

1.0

2.0

95.0

95.0

90

92

94.5

94.5

89

91

94.0

94.0

88

93.5

93.5

10 2

101

100

06

10 3

102

10 ~1

93.0

10 2

101

109

93.0

10 3

102

10 ~1

Learning rate

Perturbation radius

Learming rate

Perturbation radius

H.3.3 ResNets

In this section, we plot averages and σ -CI from 2 independent runs.

ResNets in µ P 2 transfer both the optimal learning rate and perturbation radius for SAM (Figure H.12), SAM-ON (Figure H.14) and elementwise ASAM (Figure H.15), as well as different alternatives of scaling the gradient norm contributions to SAM's denominator (Figure H.18). This suggests correctness of the derived scalings. At width multipliers 2 and 4 , µ P 2 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, ρ remains quite stable in µ P 2 but surprisingly also in SP-NAIVE. In µ P, 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 see whether even larger width would lead to suboptimal performance of µ P-global. µ P 2 is most robust to the choice of ρ and achieves the best test accuracy.

(a) SP, no

ρ

scaling

10

-

2

10

-

1

10

0

Learning rate

95

96

97

98

99

100

Training accuracy

10

-

2

10

-

1

10

0

Perturbation radius

95

96

97

98

99

100

Training accuracy

10

-

2

10

-

1

10

0

Learning rate

91

92

93

94

95

96

97

98

Test accuracy

10

-

2

10

-

1

10

0

Perturbation radius

91

92

93

94

95

96

97

98

Test accuracy

width

0.5

1.0

2.0

4.0

(b)

µ

P

2

100

100

Training accuracy

99

oenooe

99

98

98

L6

97

96

96

95

102

10 1

100

95

103

10 2

101

Leaming rate

Perturbation radius

98

98

97

97

e 1

96

96

95

95

94

94

width

93

93

0.5

1.0

92

92

2.0

4.0

91

10 ~2

101

100

91

10 3

102

10~1

Learning rate

Perturbation radius

10

-

3

10

-

2

10

-

1

Perturbation radius

ρ

93

94

95

96

97

Test accuracy

SP-naive

10

-

3

10

-

2

10

-

1

10

0

Perturbation radius

ρ

93

94

95

96

97

µ

P-naive

width multiplier

4.0

2.0

1.0

0.5

10

-

2

10

-

1

10

0

Perturbation radius

ρ

93

94

95

96

97

µ

P-global

10

-

2

10

-

1

10

0

Perturbation radius

ρ

93

94

95

96

97

µ

P

2

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 µ P 2 transfer the optimal ( η, ρ ) .

For the examples of elementwise ASAM and SAM-ON the global perturbation scaling n 1 / 2 suffices to reach µ P 2 . The stability of the optimal perturbation radius in the applied scaling n 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 a discussion of the remarkable stability of ResNets in SP. For the example of elementwise ASAM in SP, the optimal perturbation radius seems to grow.

For layerwise ASAM (Figure H.16), the optimal perturbation radius seems to grow in both SP and µ P 2 , suggesting that our scaling condition does not perfectly apply to this variant, although µ P 2 ( 97 . 09 ± 0 . 03 (+0 . 83) ) still outperforms SP ( 96 . 86 ± 0 . 05 (+0 . 83) ) in terms of the optimal test

accuracy. As Frobenius norms of weights are the only component that is not representable as a NE ⊗ OR ⊤ program, these Frobenius norms appear to scale differently than heuristically predicted over the course of training.

(a) SP, no

ρ

scaling

10

-

2

10

-

1

10

0

Learning rate

95

96

97

98

99

100

Training accuracy

10

-

1

10

0

10

1

Perturbation radius

95

96

97

98

99

100

Training accuracy

10

-

2

10

-

1

10

0

Learning rate

91

92

93

94

95

96

97

98

Test accuracy

10

-

1

10

0

10

1

Perturbation radius

91

92

93

94

95

96

97

98

Test accuracy

width

0.5

1.0

2.0

4.0

(b)

µ

P

2

100

100

Training accuracy

99

g accuracy

99

98

98

97

97

96

96

95

102

10 1

100

95

10~1

100

101

Leaming rate

Perturbation radius

98

98

width

97

97

0.5

95

96

Test accuracy

96

95

1.0

2.0

4.0

94

94

93

93

92

92

91

10 ~2

101

100

91

101

100

101

Learning rate

Perturbation radius

(a) SP, no

ρ

scaling

10

-

2

10

-

1

10

0

Learning rate

95

96

97

98

99

100

Training accuracy

10

-

1

10

0

10

1

Perturbation radius

95

96

97

98

99

100

Training accuracy

10

-

2

10

-

1

10

0

Learning rate

91

92

93

94

95

96

97

98

Test accuracy

10

-

1

10

0

10

1

Perturbation radius

91

92

93

94

95

96

97

98

Test accuracy

width

0.5

1.0

2.0

4.0

(b)

µ

P

2

100

100

Training accuracy

99

Training accuracy

99

98

98

97

97

96

96

95

10 2

10 1

109

95

10 1

109

10l

Learning rate

Perturbation radius

98

98

97

97

96

onoe

96

96

95

94

94

width

93

6

0.5

1.0

92

92

2.0

4.0

91

102

101

109

91

10 ~1

100

101

Learning rate

Perturbation radius

(a) SP, no

ρ

scaling

10

-

2

10

-

1

10

0

Learning rate

95

96

97

98

99

100

Training accuracy

10

-

3

10

-

2

10

-

1

Perturbation radius

95

96

97

98

99

100

Training accuracy

10

-

2

10

-

1

10

0

Learning rate

91

92

93

94

95

96

97

98

Test accuracy

10

-

3

10

-

2

10

-

1

Perturbation radius

91

92

93

94

95

96

97

98

Test accuracy

width

0.5

1.0

2.0

4.0

(b)

µ

P

2

100

100

Training accuracy

99

Training accuracy

99

86

L6

L6

96

96

95

10 3

10 ~2

101

100

95

10 3

102

10 ~1

Learming rate

Perturbation radius

98

98

97

97

Kemsse s1

96

Aenose s

96

95

95

94

94

width

93

93

0.5

1.0

92

92

2.0

4.0

91

103

10 ~2

101

100

91

103

10 2

10 ~1

Learning rate

Perturbation radius

H.4 Gradient norm contributions have negligible effects on generalization performance

In this section we provide ablations concerning the question which layers should contribute nonvanishingly 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 to 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 a significant effect on the test loss. This raises the question which effect the gradient normalization has in µ P. Does it contribute a scaling correction in SP, but may be dropped entirely in µ P?

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

47

47

47

47

47

46

44

18

48

48

48

48

48

48

46

17

49

49

49

49

49

50

48

16

50

50

50

50

50

50

49

16

50

50

50

50

50

50

50

23

50

49

50

50

49

50

50

28

48

48

48

49

49

48

49

39

46

46

46

46

47

46

46

37

42

42

42

43

41

41

40

22

width=64

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

51

51

51

51

51

51

48

26

52

52

52

52

52

53

50

25

52

53

53

53

53

54

52

24

53

53

53

53

54

54

54

23

53

53

53

53

53

54

55

34

52

52

52

53

52

53

55

42

51

51

51

51

52

52

54

45

50

49

49

50

50

51

51

43

44

44

45

45

45

44

46

35

width=256

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

53

53

53

53

54

53

50

30

53

54

54

54

54

55

52

30

54

54

54

54

55

56

54

29

54

54

54

54

55

56

56

29

54

55

54

54

55

57

57

39

55

55

55

55

55

56

57

46

54

54

54

54

54

55

57

49

52

52

52

52

52

52

55

48

46

46

46

47

46

48

49

40

width=1024

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

54

54

54

54

54

54

50

33

54

54

55

55

55

56

53

33

54

54

54

55

55

57

55

33

55

55

55

55

55

57

57

32

57

56

56

56

56

58

58

37

58

57

58

57

57

57

59

46

56

57

56

57

56

57

59

50

53

54

54

54

54

55

57

50

47

47

48

47

48

50

52

45

width=4096

46

47

48

49

50

Test accuracy

50

51

52

53

54

Test accuracy

53

54

55

56

57

Test accuracy

54

55

56

57

58

Test accuracy

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

46

46

46

46

46

45

43

20

48

48

48

48

49

48

45

18

49

49

49

49

50

49

47

15

50

50

50

50

50

50

48

17

50

49

50

49

50

50

49

21

49

49

49

49

49

49

49

31

48

49

48

48

48

48

48

39

47

47

46

46

46

47

47

31

42

42

41

43

43

41

40

18

width=64

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

51

51

51

51

52

51

47

23

52

52

52

52

53

53

49

21

52

53

53

53

53

54

51

19

53

53

53

53

54

55

53

18

52

53

52

53

53

55

55

28

52

51

52

53

53

53

54

41

51

51

51

51

52

52

54

47

50

50

49

50

50

50

51

43

45

46

44

45

45

45

46

25

width=256

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

53

53

53

53

53

53

48

30

53

54

54

54

55

55

51

30

54

54

54

55

56

56

53

28

54

55

55

55

56

57

55

27

54

54

55

55

56

57

57

31

56

55

55

55

55

56

58

43

54

53

53

54

53

55

57

49

52

51

52

51

52

52

55

48

46

45

46

46

48

48

49

28

width=1024

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

53

54

54

54

54

53

49

32

54

54

54

55

55

56

51

32

54

54

54

55

56

57

54

32

55

55

55

55

56

58

56

30

56

56

56

56

57

58

58

32

58

58

58

57

57

58

59

45

56

56

57

56

56

58

59

49

54

54

54

54

54

55

57

50

47

47

45

47

48

50

52

37

width=4096

46

47

48

49

50

Test accuracy

50

51

52

53

54

Test accuracy

53

54

55

56

57

Test accuracy

54

55

56

57

58

Test accuracy

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

46

46

46

46

46

46

43

20

48

48

48

48

48

48

45

18

49

49

49

49

50

49

47

16

50

50

50

50

50

50

49

18

50

50

50

50

50

50

49

24

49

49

49

49

49

50

50

33

48

49

48

47

48

49

49

39

47

46

47

46

46

47

47

41

42

42

42

41

42

42

41

29

width=64

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

51

51

51

51

52

51

48

23

52

52

52

52

52

53

50

22

52

53

53

53

53

54

52

19

53

53

53

53

53

54

54

25

52

52

53

52

53

54

55

33

52

52

52

52

52

54

55

43

51

52

51

51

52

52

54

47

50

50

50

50

49

50

52

46

45

44

46

44

44

45

46

31

width=256

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

53

53

53

53

53

53

49

31

53

54

54

54

54

55

52

31

54

54

54

55

55

56

54

29

54

55

55

55

55

57

56

27

54

55

54

55

56

57

58

33

56

55

55

54

55

56

57

46

54

54

53

55

53

55

57

49

52

51

51

51

51

53

55

45

46

46

46

46

47

48

49

37

width=1024

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

53

54

54

54

54

54

49

32

54

54

55

54

55

56

52

33

54

54

54

55

56

56

55

32

55

55

55

55

56

57

57

30

56

56

56

56

56

58

59

39

58

58

57

57

57

58

59

48

56

56

57

57

57

58

59

49

54

53

54

55

54

56

58

51

47

46

47

47

48

51

53

42

width=4096

46

47

48

49

50

Test accuracy

50

51

52

53

54

Test accuracy

53

54

55

56

57

Test accuracy

54

55

56

57

58

Test accuracy

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

46

46

46

46

46

46

44

21

48

48

48

48

48

48

45

19

49

49

49

49

49

50

48

17

50

50

50

50

50

50

49

16

50

50

49

50

50

50

50

16

49

49

49

48

48

49

50

23

48

48

48

48

48

48

48

42

47

47

47

46

47

47

47

42

42

41

42

41

41

41

42

33

width=64

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

51

51

51

51

52

51

48

24

52

52

52

52

52

53

50

23

52

52

53

53

53

54

52

21

53

52

53

53

54

54

54

22

52

52

52

53

53

54

55

37

52

52

52

53

53

54

55

46

51

51

51

51

52

52

54

49

50

51

49

50

49

50

51

48

45

46

44

46

45

45

46

38

width=256

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

53

53

53

53

53

53

50

31

53

54

54

54

54

55

52

31

54

54

54

54

55

56

54

30

54

54

55

55

56

56

56

29

54

55

55

55

55

57

58

34

56

55

56

55

55

56

58

47

54

53

54

54

54

54

57

51

52

51

51

52

52

52

54

52

46

47

47

46

47

47

49

45

width=1024

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

53

54

54

54

54

54

50

33

54

54

54

54

55

55

52

33

54

54

55

54

55

56

55

33

55

55

55

55

56

57

57

32

56

56

57

56

57

57

59

34

58

59

58

57

57

58

59

47

56

56

56

57

56

57

59

52

54

54

54

54

53

55

57

53

47

48

47

48

47

49

52

49

width=4096

46

47

48

49

50

Test accuracy

51

52

53

54

55

Test accuracy

53

54

55

56

57

Test accuracy

55

56

57

58

59

Test accuracy

For ResNets, Figure H.18(a) shows accuracies when rescaling all layers' gradient norms to Θ(1) , and Figure H.18(b) shows the results when using the original global gradient norm rescaled to Θ(1) . 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.12, optimal performance is similar but rescaling all layers' gradient norm contributions to Θ(1) may even produce a slightly more stable hyperparameter-loss landscape for ResNets.

10

-

2

10

-

1

10

0

Learning rate

95

96

97

98

99

100

Training accuracy

10

-

2

10

-

1

10

0

Perturbation radius

95

96

97

98

99

100

Training accuracy

10

-

2

10

-

1

10

0

Learning rate

91

92

93

94

95

96

97

98

Test accuracy

10

-

2

10

-

1

10

0

Perturbation radius

91

92

93

94

95

96

97

98

Test accuracy

width

0.5

1.0

2.0

4.0

(a) Rescaling all of SAM's denominator terms to

Θ(1)

10

-

2

10

-

1

10

0

Learning rate

95

96

97

98

99

100

Training accuracy

10

-

2

10

-

1

10

0

Perturbation radius

95

96

97

98

99

100

Training accuracy

10

-

2

10

-

1

10

0

Learning rate

91

92

93

94

95

96

97

98

Test accuracy

10

-

2

10

-

1

10

0

Perturbation radius

91

92

93

94

95

96

97

98

Test accuracy

width

0.5

1.0

2.0

4.0

(b) Global

∥∇

L

∥

scaling

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

In SP, we consider a global constant ρ l = ρ , whereas for µ P 2 the spectral condition ( ∗ ) requires ρ l = ρ · √ fan_out / fan_in .

Overall, SAM without layer coupling performs decently, but is outperformed by the original SAM in particular in ResNets, in µ P 2 and at large width. But note that for ResNets we adopt the hyperparameters tuned for the original SAM with layer coupling, so that these ablations only serve as preliminary experiments.

MLPs. SAM without layer coupling achieves similar optimal generalization in µ P 2 at each width compared to Figure H.9. The regime of stable ( η, ρ ) stays width-independent, but does not transfer the optimum consistently. This suggests complex or noisy dependence of the training dynamics on ρ .

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

46

46

46

46

46

45

42

12

48

48

48

48

48

47

44

11

49

49

49

49

49

48

45

11

49

50

50

50

49

49

46

11

50

50

50

50

49

49

46

11

49

50

49

49

49

49

11

11

49

49

48

48

48

47

11

11

46

47

47

48

46

46

11

11

42

41

41

41

42

40

11

11

width=64

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

51

51

51

51

51

51

48

13

51

52

52

52

52

52

51

12

52

52

52

53

53

53

51

11

53

52

53

53

53

53

52

11

53

53

53

53

53

54

52

11

52

52

51

52

51

53

52

11

51

52

51

51

51

51

50

13

49

49

48

49

51

49

47

12

45

45

44

45

44

44

11

12

width=256

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

53

53

53

53

54

54

51

18

53

53

53

54

54

54

53

15

53

54

54

54

54

55

54

13

54

54

54

55

55

55

55

12

55

54

55

54

55

55

55

13

55

55

56

56

56

54

55

11

53

54

54

54

54

54

54

12

51

51

51

51

52

53

50

12

46

47

48

46

47

47

11

11

width=1024

0.0

0.001

0.0032

0.01

0.0316

0.1

0.3162

1.0

Perturbation radius

0.0283

0.0476

0.0801

0.1346

0.2264

0.3808

0.6404

1.0771

1.8114

Learning rate

53

54

54

54

54

54

52

23

54

54

54

54

55

55

54

20

54

54

54

54

55

56

55

17

54

55

55

55

55

56

56

15

56

56

56

56

56

56

56

14

58

59

57

58

57

57

57

13

56

56

56

56

56

57

57

13

53

53

53

53

53

55

53

13

47

48

48

47

48

48

14

14

width=4096

46

47

48

49

50

Test accuracy

49

50

51

52

53

Test accuracy

51

52

53

54

55

Test accuracy

54

55

56

57

58

Test accuracy

(a)

µ

P

2

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 SP and µ P 2 , in particular at large width. As expected, ρ transfers in µ P 2 .

10

-

3

10

-

2

10

-

1

Perturbation radius

93

94

95

96

97

98

Test accuracy

width

4.0

2.0

1.0

0.5

(a) SP

10

-

3

10

-

2

10

-

1

Perturbation radius

93

94

95

96

97

98

Test accuracy

width

4.0

2.0

1.0

0.5

(b)

µ

P

2

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. µ P 2 always achieves the best final test accuracy. In 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 µ P 2 uniformly dominates over the entire course of training. Our theory suggests that in µ P 2 the gradients are 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 ViTs, µ P generally achieves decent accuracy faster than SP, since gradient norms are already scaled correctly at 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 to SAM in µ P 2 at the end of training.

0

25

50

75

100

125

150

175

200

Epochs of Training

10

1

10

2

100% - test accuracy

SAM

µ

P

2

SGD

µ

P

SAM SP

SGD SP

(a) ResNet on CIFAR10

0

50

100

150

200

250

300

Epochs of training

3

×

10

1

4

×

10

1

6

×

10

1

100% - test accuracy

SAM

µ

P

2

AdamW

µ

P

SAM SP

AdamW SP

(b) ViT on CIFAR100

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]

Guidelines:

Justification: In the abstract and introduction we state our main contributions while acknowledging related work. All main claims are theoretically proven and/or empirically verified.

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]

2. Limitations

Justification: Limitations are discussed in the future work section as well as in the section in the appendix that is related to the respective limitation.

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.

Guidelines:

The authors are encouraged to create a separate "Limitations" section in their paper.

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 a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The paper should point out 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 factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

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 important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

3. Theory Assumptions and Proofs

Answer: [Yes]

Guidelines:

Justification: We state all assumptions in the main paper and Appendix C, and provide all formal proofs in Appendix E.

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

All assumptions should be clearly stated or referenced in the statement of any theorems.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

4. Experimental Result Reproducibility

Theorems and Lemmas that the proof relies upon should be properly referenced.

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

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.

Answer: [Yes]

Guidelines:

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.

The answer NA means that the paper does not include experiments.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

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 a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

5. Open access to data and code

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

contains the essential modifications.We are working on making Python code to reproduce all of our experiments publicly available. Guidelines:

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

The answer NA means that paper does not include experiments requiring code.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The 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 instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Justification: All experimental details are disclosed in the main paper or Appendix G.

Answer: [Yes]

Guidelines:

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The answer NA means that the paper does not include experiments.

The full details can be provided either with the code, in appendix, or as supplemental material.

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

7. Experiment Statistical Significance

Answer: [Yes]

The answer NA means that the paper does not include experiments.

Justification: As stated in Appendix G, we repeat all main experiments with multiple independent runs and report confidence bands within the empirical 2 . 5% - to 97 . 5% -quantiles. When we repeat experiments on Vision Transformers that we have also conducted on MLPs or ResNets, we do not use multiple runs due to limitations in computational resources. Guidelines:

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

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, or overall run with given experimental conditions).

The assumptions made should be given (e.g., Normally distributed errors).

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 a 96% CI, if the hypothesis of Normality of errors is not verified.

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

8. Experiments Compute Resources

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.

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Justification: We provide the type of GPU used and number of GPU seconds required for each experiment in Appendix G.

Answer: [Yes]

Guidelines:

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The answer NA means that the paper does not include experiments.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

9. Code Of Ethics

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Guidelines:

Justification: We provide a theoretical analysis of a widely used optimization algorithm, point out the algorithm's limitations in large models and propose a correction. We do not foresee any ethical concerns.

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

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.

Answer: [NA]

Guidelines:

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

The answer NA means that there is no societal impact of the work performed.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The 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

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 to train models that generate Deepfakes faster.

technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

11. Safeguards

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Justification: We only use standard vision architectures and vision datasets in our experiments and do not release any data or models.

Answer: [NA]

Guidelines:

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.

The answer NA means that the paper poses no such risks.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

12. Licenses for existing assets

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

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]

The answer NA means that the paper does not use existing assets.

Justification: We cite the standard CIFAR10, CIFAR100 (Krizhevsky et al., 2009) and ImageNet1K (Deng et al., 2009) datasets following the standard practice. We also cite the Python assets PyTorch (Paszke et al., 2019), mup (Yang et al., 2022) and the GitHub repository implementing SAM (Samuel, 2022) that we use as a basis for our experiments. Guidelines:

The authors should cite the original paper that produced the code package or dataset.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

The authors should state which version of the asset is used and, if possible, include a URL.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

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 assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

13. New Assets

Justification: The paper does not release new assets.

Answer: [NA]

Guidelines:

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The answer NA means that the paper does not release new assets.

The paper should discuss whether and how consent was obtained from people whose asset is used.

14. Crowdsourcing and Research with Human Subjects

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

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

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

Answer: [NA]

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

Answer: [NA]

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

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.

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.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- The idea is original, significant and seems to work as shown by the experiments\n\nWEAKNESSES:\nTo 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe authors somewhat address the limitations in the last section.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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).\n\nQUESTIONS:\n- 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.\"\n\nLIMITATIONS:\nYes, the authors adequately addressed the limitations", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\nSome 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?\n\nLIMITATIONS:\nThe authors do include some discussion of some limitations in the future work section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}

comments: 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.

STRENGTHS:
- This paper is well-written.
- 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].

[1] Efficient Sharpness-aware Minimization for Improved Training of Neural Networks.

[2] Surrogate Gap Minimization Improves Sharpness-Aware Training.

QUESTIONS:
N/A

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-27