Critical Influence of Overparameterization on Sharpness-aware Minimization

## Abstract

Abstract Training an overparameterized neural network can yield minimizers of different generalization capabilities despite the same level of training loss.
Meanwhile, with evidence that suggests a strong correlation between the sharpness of minima and their generalization errors, increasing efforts have been made to develop optimization methods to explicitly find flat minima as more generalizable solutions.
Despite its contemporary relevance to overparameterization, however, this sharpness-aware minimization (SAM) strategy has not been studied much yet as to exactly how it is affected by overparameterization.
Hence, in this work, we analyze SAM under overparameterization of varying degrees and present both empirical and theoretical results that indicate a critical influence of overparameterization on SAM.
At first, we conduct extensive numerical experiments across vision, language, graph, and reinforcement learning domains and show that SAM consistently improves with overparameterization.
Next, we attribute this phenomenon to the interplay between the enlarged solution space and increased implicit bias from overparameterization.
Further, we prove multiple theoretical benefits of overparameterization for SAM to attain (i) minima with more uniform Hessian moments compared to SGD, (ii) much faster convergence at a linear rate, and (iii) lower test error for two-layer networks.
Last but not least, we discover that the effect of overparameterization is more significantly pronounced in practical settings of label noise and sparsity, and yet, sufficient regularization is necessary.

## 1 Introduction

The recent success of deep learning can be attributed to the increasing size of neural networks: as the networks become deeper and wider, they achieve superior performance across various applications (Kaplan et al., ; Alayrac et al., ; Dehghani et al., ; Radford et al., ).
How come such large—and thus potentially overparameterized—neural networks work so well?
While a complete understanding remains elusive, research suggests that overparameterization can help multiple aspects of the learning process including even generalization (Neyshabur et al., ; Du and Lee, ).
In particular, overparameterized neural networks show convexity-like behavior during optimization, making all local minima globally optimal, and thereby, enabling one to easily find global minima with local optimization methods such as gradient descent (Kawaguchi, ; Du et al., ).

However, not all global minima are necessarily equivalent;
converging to different minima can yield a large disparity of generalization capabilities, despite their same level of training loss reaching almost zero (Keskar et al., ; Gunasekar et al., ).
One plausible explanation for such an implicit phenomenon is that generalization is somewhat negatively correlated with the sharpness of the loss landscape, *i.e*., flat minima tend to generalize better than sharp ones (Chaudhari et al., ; Jiang et al., ).
This calls for new ways to guide the optimization process to converge to flat minima, and various strategies have been suggested to this end (Izmailov et al., ; Foret et al., ; Orvieto et al., ).
Indeed, it has been observed in many cases that minimizing sharpness improves generalization performance (Bahri et al., ; Chen et al., ; Qu et al., ).

Despite the sharpness minimization scheme being developed in the context of overparameterization, the precise impact of overparameterization on the scheme has not been studied much in the literature.
Hence, we analyze the effects of overparameterization on sharpness-aware minimization, or SAM (Foret et al., ), to various extents in this work.
More specifically, we conduct a wide range of experiments and precisely measure the effects of overparameterization, ranging from standard tasks in computer vision and natural language processing, to molecular property prediction, and further, to video game in reinforcement learning.
We then delve into the results through detailed analyses of the solution space and optimization trajectories on a simple regression problem as well as the relationship between overparameterization and implicit bias of SAM.
Also, we explore the theoretical implications of overparameterization on SAM including the characteristics of the attained minima, convergence rate, and generalization properties.
In addition, we investigate how the effects of overparameterization change under different scenarios including label noise, sparsity, and regularization.

Our key findings are summarized as follows.

- •
In Section 3, our extensive evaluations reveal a consistent trend that overparameterization improves the generalization benefit of SAM, and SAM may not take effect without it.
- •
In Section 4, we show that the critical influence of overparameterization is due to the enlarged solution space and increased implicit bias of SAM under overparameterization.
- •
In Section 5, we demonstrate some theoretical benefits of overparameterization for SAM to (i) attain minima with more uniform Hessian moments compared to SGD, (ii) achieve much faster convergence at a linear rate, and (iii) attain lower test error for two-layer ReLU networks.
- •
In Section 6, we discover several merits and caveats of overparameterization in that the benefits of overparameterization can be more pronounced under label noise or sparsity while sufficient regularization is required to attain the benefit.

## 2 Background

#### Sharpness-aware minimization

Consider the general unconstrained optimization problem:

$$ $\min_{x}f(x)$ (1) $$

where $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ is the objective function to minimize, and $x\in\mathbb{R}^{d}$ is the optimization variable.
Based on recent observations that indicate a correlation between the sharpness of $f$ at a minimum and its generalization error (Keskar et al., ; Dziugaite and Roy, ; Jiang et al., ), Foret et al. () suggest to turn (1) into a min-max problem of the following form

$$ $\min_{x}\max_{\lVert\epsilon\rVert_{2}\leq\rho}f(x+\epsilon)$ (2) $$

where $\epsilon$ and $\rho$ denote some perturbation added to $x$ and its norm, respectively.
Thus, the goal is now to seek $x$ that minimizes $f$ in its entire $\epsilon$-neighborhood, such that the objective landscape becomes locally flat.
Taking the first-order Talyor approximation of $f$ at $x$ and solving for optimal $\epsilon^{\star}$ gives the following update rule for SAM:

$$ $x_{t+1}=x_{t}-\eta\nabla f\left(x_{t}+\rho\frac{\nabla f(x_{t})}{\lVert\nabla f (x_{t})\rVert_{2}}\right).$ (3) $$

#### Related work

The effectiveness of SAM on improving generalization has been verified in the literature across different settings (Chen et al., ; Kaddour et al., ; Bahri et al., ).
Subsequent works have analyzed various aspects of SAM including convergence rate (Andriushchenko and Flammarion, ; Mi et al., ; Si and Yun, ) and implicit bias (Compagnoni et al., ; Wen et al., ; Andriushchenko et al., ).
Our work takes a closer look at the effect of overparameterization on SAM, which was left unattended in the literature despite its importance.

Meanwhile, a considerable amount of evidence has indicated the benefit of overparameterization for training neural networks from multiple perspectives.
On the empirical side, overparameterization leads to better performance across various domains (Kaplan et al., ; Radford et al., ; Dehghani et al., ).
On the theoretical side, it turns all local minima into global ones and thus allows one to easily find global minima despite the non-convex objective (Kawaguchi, ; Du et al., ).
Recent literature has also proved the power of overparameterization to enable much faster convergence (Ma et al., ; Vaswani et al., ; Meng et al., ) and better generalization (Allen-Zhu et al., ; Brutzkus and Globerson, ).
While the previous work is mostly focused on the effect of overparameterization for non-sharpness-aware optimizers, our work studies its effect on SAM through in-depth analyses and comprehensive evaluations.

## 3 Key observation: SAM improves with overparameterization

**Table 1: Summary of evaluation workloads. They cover eight different datasets spanning five domains and six tasks at varying scale, and include eight neural network models of five different architecture types. For each workload, we test up to ten different models of varying degrees of parameterization. We defer to Section A.1 for the full experiment details.**
| Workload # | Domain | Task | Dataset | Architecture | Model |
| --- | --- | --- | --- | --- | --- |
| $1$ | Synthetic | Regression | Synthetic | MLP | Two-layer MLP |
| $2$ | Vision | Image classification | MNIST | MLP | LeNet-300-100 |
| $3$ | Vision | Image classification | CIFAR-10 | CNN | ResNet-18 |
| $4$ | Vision | Image classification | ImageNet | CNN | ResNet-50 |
| $5$ | Language | PoS tagging | Universal Dependencies | Transformer | Encoder-only Transformer |
| $6$ | Language | Sentiment classification | SST-2 | RNN | LSTM |
| $7$ | Chemistry | Graph property prediction | ogbg-molpcba | GNN | GCN |
| $8$ | Game | Proximal policy optimization | Atari Breakout | CNN | Five-layer CNN |

Figure: Figure 1: Improvement in validation metrics by SAM. The generalization benefit of SAM tends to increase as the model becomes more overparameterized. We present the full results including the absolute metrics for SAM and baseline optimizers in Figure 7 of Appendix B.
Refer to caption: x1.png

SAM is introduced to find flat minima and thereby improve generalization performance in practice.
In this work, we are interested in whether and how this improvement is affected by overparameterization.
In order to understand any potential relationship between SAM and overparameterization, we first focus on precisely measuring the effect of overparameterization.
More specifically, we conduct a wide range of deep learning experiments (see Table 1 for the summary of all tested workloads), and observe how the generalization improvement made by SAM changes with respect to overparameterization.

As a result, we find a strong and consistent trend that SAM improves with overparameterization in all tested cases (see Figure 1).
To elaborate, initially SAM does *not* work much better than the non-sharpness-aware baseline optimizer when the model is at relatively low number of parameters, and yet, it starts to improve with more parameters and makes a clear distinction at very large number of parameters.
We emphasize that this holds true for a variety of architectures (MLP, CNN, RNN, GCN, Transformer) and datasets of different domains (Synthetic, Vision, Language, Chemistry, Game).
This result possibly indicates that SAM is more effective, if not only, when applied to overparameterized models.
We note that some evidence of the similar positive influence of overparameterization for SAM can be derived in the literature (Chen et al., ), however, no prior work has conducted experiments or confirmed this phenomenon at any scale comparable to ours.

## 4 Understanding why SAM improves with overparameterization

How does SAM benefit from overparameterization? Why does overparameterization particularly favor SAM over non-sharpness-aware optimizers?
We address these questions in this section to develop a better understanding of the effect of overparameterization on SAM.
Precisely, we posit that it is potentially due to the complementarity between overparameterization enlarging the solution space and the implicit bias of SAM driving to reach flat minima;
*i.e*., once there are more diverse solutions available including both sharp and flat minima (by overparameterization), optimizers biased toward flat solutions (such as SAM) will more likely find such solutions than unbiased optimizers (such as GD).
We support this reasonable hypothesis by demonstrating the following: (1) SAM does find simpler and flatter solutions than GD, and (2) the implicit bias of SAM is amplified with increasing optimal perturbation bound, both when the model is overparameterized.

### 4.1 Overparameterization leads to simpler and flatter solutions for SAM

To corroborate our hypothesis, we start with a simple experiment where we train one-hidden-layer ReLU networks using full-batch SAM and GD following Andriushchenko and Flammarion ();
we use $10$ and $100$ hidden neurons for underparameterized and overparameterized cases, respectively;
we run three random seeds and compare solutions obtained by SAM and GD.
The results are plotted in Figure 2.

First, we find that the solutions found by SAM are not differentiated much from those of GD when the model has only 10 neurons (see Figures 2(a) and 2(b)).
Looking closely, they all seem to be roughly $4$ to $6$ degrees of piecewise linear functions, *i.e*., the number of line segments for each solution is less than $10$, which is the maximum possible joints that this model can have.
On the other hand, in the case of $100$ neurons, one can easily see that the solutions found by SAM are much simpler (and thus more likely to generalize) compared to those by GD (see Figures 2(c) and 2(d)).

Next, we also track the optimization trajectories of both SAM and GD.
The trajectories are plotted along PCA directions calculated from the converged minima following Li et al. ().
The results are illustrated in Figure 3(a).
We find that both SAM and GD reach solutions in a similar basin when the model is underparameterized, whereas in the overparameterized case, they reach different solutions, *i.e*., SAM reaches a flatter solution, even though they all start from the same initial point.

These results indeed support the idea that SAM has some implicit bias that drives itself toward a certain type of solutions (*e.g*., simple and flat) as previously shown in prior work (Andriushchenko and Flammarion, ; Compagnoni et al., ; Wen et al., ).
More importantly, however, these results reveal that overparameterization is a critical factor in facilitating this implicit behavior of SAM, without which the space of potential solutions decreases, and SAM may not take effect.

Figure: (a)
Refer to caption: x9.png

Figure: (a) Optimization trajectories
Refer to caption: x13.png

### 4.2 Implicit bias of SAM increases with overparameterization

Figure: Figure 4: Validation accuracy versus $\rho$ for ResNet-18 and CIFAR-10. $\rho^{\star}$ is located to be higher with more parameters. See Figures 8, 9, 10 and 11 of Appendix B for more results on other workloads.
Refer to caption: x17.png

While overparameterization can secure favorable conditions for SAM, it is not to be confused with guaranteeing the implicit bias of SAM taking effect.
In fact, we can relate the implicit bias of SAM to the perturbation bound $\rho$ to bridge this gap.
Specifically, SAM can be proved to be an SDE model of SGD on an implicitly regularized loss (Compagnoni et al., ):

$$ $\tilde{f}(x)\coloneqq f(x)+\rho\mathbb{E}\|\nabla f_{\gamma}(x)\|_{2}$ (4) $$

where $\tilde{f}$ denotes an SGD approximation of SAM objective and $\gamma$ refers to some stochasticity.
This indicates that SAM becomes more regularized (*i.e*., the implicit bias is amplified) when $\rho$ increases.(^1^11This holds as long as $\rho$ is not too large, by which it might overshadow minimizing $f$ and implicitly bias the optimizer toward stationary points such as saddles and maxima. Note it reduces to standard SGD when $\rho=0$.)

We verify this by empirically finding the optimal perturbation bound $\rho^{\star}$ that yields the best generalization performance as we change the degree of overparameterization.
Specifically, we take a standard deep learning task and perform an extensive grid search to find $\rho^{\star}$.
The result is displayed in Figure 4.
It is observed that $\rho^{\star}$ tends to increase as the number of parameters increases;
*i.e*., seeing from left to right, $\rho$ value that yields highest accuracy (marked as green star $\star$) tends to increase.
This result is certainly encouraging since it supports that the generalization benefit of SAM obtained via implicit regularization can indeed increase by overparameterization.

Additionally, we can develop a conceptual account of why $\rho$ increases with overparameterization:
the scale of perturbation $\epsilon$ controlled by $\rho$ needs to increase relatively proportionally with respect to the increase in model dimensionality in order to preserve a similar level of perturbation effect;
an extreme case is when the model grows infinitely large, $\rho$ should grow together, otherwise no perturbation can be made, *i.e*., $\|\epsilon\|\rightarrow\infty$ as $d\rightarrow\infty$.

A non-asymptotic interpretation can also be developed from an optimization perspective.
A bound from Lipschitz smoothness: $\left\lVert\nabla f\left(x+\rho\frac{\nabla f(x)}{\|\nabla f(x)\|_{2}}\right)-
\nabla f(x)\right\rVert_{2}\leq\beta\left\lVert x+\rho\frac{\nabla f(x)}{\|
\nabla f(x)\|_{2}}-x\right\rVert_{2}=\beta\rho$ indicates that the SAM gradient becomes more similar to the original gradient as the model gets smoother (*i.e*., smaller smoothness constant $\beta$) with increasing size , requiring larger perturbation bound to achieve similar levels of perturbation effect.
These hold under the assumption that overparameterization makes the model smoother, which we empirically confirm in Figure 18(a) of Appendix D.

## 5 Theoretical benefits of overparameterization for SAM

Figure: (a) Non-uniformity
Refer to caption: x23.png

Our understanding of overparameterization has been significantly advanced in recent years through a variety of theoretical analyses such as convergence (Ma et al., ; Vaswani et al., ), generalization (Allen-Zhu et al., ; Brutzkus and Globerson, ), and implicit bias (Neyshabur, ; Zhang et al., ).
Despite this, there is little work on its implications for SAM.
In this section, we present theoretical analyses of the effect of overparameterization on improving SAM from multiple different angles.
Specifically, we discover that (1) linearly stable minima for SAM have more uniform Hessian moments compared to SGD, (2) SAM can converge much faster, and (3) achieve lower test error, all when the model is overparameterized.

### 5.1 Linearly stable minima of SAM have a more uniform Hessian compared to SGD

It has been observed in Woodworth et al. (); Xie et al. () that SGD converges to certain types of minima among many others in an overparameterized regime.
We theoretically analyze the minima attained by SAM and how they compare to the minima attained by SGD from the perspective of linear stability (Wu et al., , ), which is defined as follows:

###### Definition 5.1 .

(Linear stability)
Consider an iterative first-order optimizer $x_{t+1}=x_{t}-\eta_{t}G(x_{t})$ where $\eta_{t}$ denotes a step size and $G$ refers to a stochastic gradient estimate measured at the current iterate $x_{t}$.
A minimizer $x^{\star}$ is called linearly stable if there exists a constant $C$ such that

$$ $\mathbb{E}[\|\tilde{x}_{t}-x^{\star}\|^{2}]\leq C\|\tilde{x}_{0}-x^{\star}\|^{2}$ $$

for all $t>0$ under $\tilde{x}_{t+1}=\tilde{x}_{t}-\nabla G(x^{\star})(\tilde{x}_{t}-x^{\star})$, *i.e*., if it does not deviate far from $x^{\star}$ once arrived near a fixed point.

With this, we provide the following theorem of a linearly stable minima for a stochastic SAM.

###### Theorem 5.2 .

Let us assume $x^{\star}=0$ without loss of generality.
Then $x^{\star}$ is linearly stable for a stochastic SAM if the following is satisfied:

$$ $\begin{split}\lambda_{\textup{max}}&\left((I-\eta H-\eta\rho H^{2})^{2}+\eta( \eta-2\rho)(M_{2}-H^{2})\right.\\ &\hskip 10.00002pt\left.+2\eta^{2}\rho(M_{3}-H^{3})+\eta^{2}\rho^{2}(M_{4}-H^{ 4})\right)\leq 1\end{split}$ (5) $$

where $H=\frac{1}{n}\sum_{i=1}^{n}H_{i}$ and $M_{k}=\frac{1}{n}\sum_{i=1}^{n}H_{i}^{k}$ are the average Hessian and the $k$-th moment of the Hessian at $x^{\star}$ over $n$ training data.
Subsequently as a necessary condition of (5) it follows that

$$ $\displaystyle\begin{split}&0\leq a(1+\rho a)\leq\frac{2}{\eta},\quad 0\leq s_{ 2}^{2}\leq\frac{1}{\eta(\eta-2\rho)},\quad 0\leq s_{3}^{3}\leq\frac{1}{2\eta^{ 2}\rho},\quad 0\leq s_{4}^{4}\leq\frac{1}{\eta^{2}\rho^{2}},\end{split}$ (6) $$

where $a=\lambda_{\text{max}}(H),s_{k}=\lambda_{\text{max}}((M_{k}-H^{k})^{1/k})$ are the sharpness and the non-uniformity of the Hessian measured with the $k$-th moment, respectively.

The detailed proof of the theorem is provided in Appendix G.

Our result (6) suggests that SAM requires less sharp minima and more uniformly distributed Hessian moments to achieve linear stability (provided that $\rho>0$) compared to those of SGD (Wu et al., ), *i.e*., when $\rho\rightarrow 0$ in (6).
While a similar result is shared by a concurrent work of Behdin et al. (), we further ensure that higher-order terms of Hessian moments are bounded, and interestingly, it becomes tighter for a larger $\rho$.
To corroborate our result, we measure the empirical sharpness and non-uniformity of Hessian.
The results are reported in Figures 3(b), 3(c) and 5(a).

###### Definition 5.1 .

###### Theorem 5.2 .

### 5.2 Stochastic SAM converges much faster with overparameterization

Prior works have revealed the power of overparameterization for different stochastic optimizers in the sense that it enables much faster converge (Ma et al., ; Vaswani et al., ; Meng et al., ).
We show that this theoretical efficacy of overparameterization also extends to a stochastic SAM.

Let us start by providing some assumptions used below.

###### Definition 5.3 .

(Smoothness) $f$ is $\beta$-smooth if there exists $\beta>0$ s.t. $\|\nabla f(x)-\nabla f(y)\|\leq\beta\|x-y\|$ for all $x,y\in\mathbb{R}^{d}$.

###### Definition 5.4 .

(Polyak-Lojasiewicz) $f$ is $\alpha$-PL if there exists $\alpha>0$ s.t. $\|\nabla f(x)\|^{2}\geq\alpha(f(x)-f(x^{\star}))$ for all $x\in\mathbb{R}^{d}$.

These assumptions are quite standard and used frequently in the optimization literature.
The smoothness assumption is satisfied for any neural network with smooth activation and loss function with bounded inputs (Andriushchenko and Flammarion, ), and the Polyak-Lojasiewicz (PL) condition is satisfied when the model is overparameterized (Belkin, ; Liu et al., ), which we empirically verify in Figure 18(b) of Appendix D.

Also, here we define overparameterization in terms of interpolation: the model possessing a sufficient number of parameters to interpolate the whole training data and achieving zero training loss for the empirical risk $f(x)\coloneq\sum_{i=1}^{n}f_{i}(x)$.
From a stochastic optimization perspective, this means that there exists a stationary point $x^{\star}$ for every data point as follows:

###### Definition 5.5 .

(Interpolation) There exists $x^{\star}$ s.t. $f_{i}(x^{\star})=0$ and $\nabla f_{i}(x^{\star})=0$ for $i=1,\ldots,n$.

Under these assumptions, we present the following convergence theorem of a stochastic SAM:

###### Theorem 5.6 .

Suppose each $f_{i}$ is $\beta$-smooth, $f$ is $\lambda$-smooth and $\alpha$-PL, and interpolation holds.
For any $\rho\leq\frac{1}{(\beta/\alpha+1/2)\beta}$, a stochastic SAM that runs for $t$ iterations with constant step size $\eta^{\star}\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{\alpha-(\beta+\alpha
/2)\beta\rho}{2\lambda\beta(\beta\rho+1)^{2}}$ gives the following convergence guarantee:

$$ $\underset{x_{t}}{\mathbb{E}}\left[f(x_{t})\right]\leq\left(1-\frac{\alpha-( \beta+\alpha/2)\beta\rho}{2}\,\eta^{\star}\right)^{t}\,f(x_{0}).$ $$

We provide the full proof in Appendix E, which also contains result for the more general case of a mini-batch SAM.

This result shows that with overparameterization, a stochastic SAM can converge as fast as the deterministic gradient method at a linear convergence rate, which is much faster than the well-known sublinear rate of $\mathcal{O}(1/t)$ for SAM (Andriushchenko and Flammarion, ).
We corroborate our result empirically as well, by measuring how training proceeds with overparameterization on realistic workloads.
The results are reported in Figure 5(b).

###### Definition 5.3 .

###### Definition 5.4 .

###### Definition 5.5 .

###### Theorem 5.6 .

### 5.3 Test error of SAM can decrease with overparameterization

Recent works have demonstrated that overparameterization can even improve generalization both empirically and theoretically (Neyshabur et al., ; Brutzkus and Globerson, ).
Here, we theoretically show that overparameterization also improves generalization for SAM in the sense that test error can decrease with larger network widths.

To this end, we follow the same setting of Allen-Zhu et al. ().
More specifically, we consider a risk minimization over some unknown data distribution $\mathcal{D}$ using a one-hidden-layer ReLU network with a smooth convex loss function (*e.g*., cross entropy).
The network is assumed to be initialized with Gaussian and take bounded inputs.
Then, we characterize a generalization property of a stochastic SAM as below.

###### Theorem 5.7 .

(Informal)
Suppose we train a network having $m$ hidden neurons with training data sampled from $\mathcal{D}$.
Then, for every $\varepsilon$ in some open interval, there exists $M\propto 1/\varepsilon$ such that for every $m\geq M$, with appropriate values of $\eta,\rho,T$, a stochastic SAM gives the following guarantee on the test loss with high probability:

$$ $\underset{x_{0},\cdots,x_{T-1}}{\mathbb{E}}\left[\frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}_{\mathcal{D}}[f(x_{t})]\right]\leq\varepsilon.$ $$

We present a formal version of the theorem and its proof in Appendix F.

This result suggests that to achieve $\varepsilon$-test accuracy from running $T$ iterations of SAM requires a minimum width $M$ proportional to $1/\varepsilon$.
This indicates that a network with a larger width can achieve a lower test error, and hence, overparameterization can improve generalization for SAM.
We support our result empirically as well on synthetic data for a simple regression task.
The results are reported in Figure 5(c).

###### Theorem 5.7 .

## 6 Further merits and caveats of overparameterization in practice

So far we have investigated the critical influence of overparameterization on SAM from various perspectives.
In this section, we explore further merits and caveats of overparameterization which could serve as a package of guidance on employing SAM in practice.
Specifically, we observe that the overparameterization benefit of SAM continues to exist and can become more evident under label noise or sparsity.
We also discover that sufficient regularization is required to attain the benefit of overparameterization for SAM.

#### Benefit of overparameterization for SAM is more pronounced under label noise.

Figure: (a) Label noise
Refer to caption: x27.png

In practice, deep learning models are frequently trained on noisy data (Song et al., ).
To further examine whether the overparameterization benefit for SAM continues to exist in this scenario, we introduce some label noise to training data (Angluin and Laird, ; Natarajan et al., ) and see how SAM responds.
The results are reported in Figure 6(a).
Overall, we find SAM benefits from overparameterization significantly more than SGD in the presence of label noise.
Precisely, the accuracy improvement made by SAM keeps on increasing as the model becomes more overparameterized, and this trend is more pronounced with a higher noise level;
*e.g*., it rises from $5\%$ to nearly $50\%$ at the highest noise rate.
We note that SAM is previously known to be more robust to label noise compared to SGD (Foret et al., ; Baek et al., ), and yet this result reveals that overparameterization has a profound influence on the robustness of SAM, which could be useful for scenarios where learning has to be done with noisy data.

#### SAM benefits from sparse overparameterization.

There has been a recent interest in employing sparsity to train large models to alleviate the computation and memory costs (Hoefler et al., ; Mishra et al., ).
To test the effect of overparameterization on SAM under this setting, we introduce a varying degree of sparsity to an overparameterized model at initialization (Lee et al., ) such that the number of parameters matches the original dense model.
The results are reported in Figure 6(b).
We observe that the generalization improvement tends to increase as the model becomes more sparsely overparameterized;
more precisely, the average accuracy improvement increases from $0.4\%$ in the small dense model to around $0.8\%$ in the large sparse model.
This result suggests that one can consider taking sparsification more actively when employing SAM.

#### Sufficient regularization is needed to secure the benefit of overparameterization for SAM.

We also investigate whether the overparameterization benefit for SAM continues to exist in environments where models are prone to overfitting due to insufficient regularization (Ying, ).
Specifically, we evaluate three cases as follows:
(a) without weight decay,
(b) without early stopping, and
(c) without sufficient inductive bias.(^2^22For (c), we train vision transformers for CIFAR-10 where the models are prone to overfitting without being pre-trained on a massive dataset due to the lack of inductive biases inherent to CNNs (Lee et al., ; Chen et al., ).)
The results are reported in Figures 6(c), 6(d) and 6(e).
We observe that the generalization improvement does not increase by simply adding more parameters.
The results indicate that securing sufficient regularization is required in practice to attain the overparameterization benefit for SAM.

## 7 Discussion and future work

In this work, we have uncovered critical influence of overparameterization on SAM from both empirical and theoretical perspectives.
We started with an extensive evaluation to discover a highly consistent trend that the generalization benefit of SAM increases with overparameterization, without which SAM may not take effect (see Section 3).
This led us to come up with a reasonable hypothesis to explain the benefit in terms of solution space and implicit bias (see Section 4).
In addition to the empirical advantages, we analyzed the overparameterization benefit for SAM from various theoretical aspects including linear stability, convergence, and generalization (see Section 5).
Finally, we explored the merits and caveats of overparameterization which can be especially useful for practitioners to fully exploit the benefit of SAM when training large-scale models (see Section 6).
We believe these findings serve as a bridge to connect overparameterization and SAM, which were explored rather separately in the literature despite the strong link between them.
We discuss the practical implications of our results and possible future research directions in more detail below.

#### Potential to modern deep learning

Our key observations in Section 3 indicate a great potential to use SAM in the modern landscape of large-scale training (Kaplan et al., ; Belkin, ).
Also, our results in Section 6 further highlight its potential in the current trend where foundation models are often trained with noisy data (Radford et al., ; Schuhmann et al., ) or to employ sparsity (Frantar et al., ; Jiang et al., ).
In this regard, we can possibly anticipate that the overparameterization benefit might hold even when training billion-scale foundation models (Zhang et al., ; Dehghani et al., ), which we leave to explore as future work.
It would also be interesting to study how popular settings for training foundation models other than label noise or sparsity affect the benefit, such as quantization (Gholami et al., ), differential privacy (Yu et al., ), or human alignment (Ouyang et al., ).

#### More ablation study

In addition to our explanation of how overparameterization improves SAM in Section 4, we have conducted an additional ablation study to investigate whether the increased benefit of SAM can be attributed to other factors.
Inspired by recent studies suggesting that overparameterized models can behave like linearized models (Jacot et al., ) while such implicit linearization phenomenon can coincide independently of overparameterization (Chizat et al., ), we have tested if the increased benefit of SAM is due to linearization.
As a result, we have observed that SAM underperforms SGD in the linearized regimes by more than $10\%$.
This indicates that linearization is not the main factor for the increased benefit of SAM and again verifies that overparameterization itself is likely to be the main factor of the benefit.
Detailed results can be found in Appendix C.

#### Extension to other sharpness minimization schemes

Our work analyzes the effect of overparameterization on SAM, which is one of the most popular sharpness minimization schemes.
Given that different sharpness minimization schemes can
make a difference in the found minima and resulting performance (Kaddour et al., ), extension of our analyses to other schemes (Izmailov et al., ; Orvieto et al., ) and studying how they compare to SAM under overparameterization would be a promising avenue for future work.
Also, recent works suggest that different versions of SAM can lead to different convergence properties (Si and Yun, ; Dai et al., ) or implicit biases (Wen et al., ; Compagnoni et al., ).
We plan to extend our theoretical analyses in Section 5, which were developed for unnormalized $n$-SAM, to different versions so as to reveal other interesting characteristics of SAM under overparameterization.

## Acknowledgement

This work was partly supported by the Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korean government (MSIT) (IITP-2019-0-01906, Artificial Intelligence Graduate School Program (POSTECH) and RS-2022-II220959, (part2) FewShot learning of Causal Inference in Vision and Language
for Decision Making), the National Research Foundation of Korea (NRF) grant funded by the Korean government (MSIT) (2022R1F1A1064569, RS-2023-00210466, RS-2023-00265444), and POSCO Creative Ideas grant (2023Q032).
Sungbin Shin was supported by Kwanjeong Educational Foundation Scholarship.
M.A. was supported by the Google Fellowship and Open Phil AI Fellowship.

## Appendix A Experimental details

### A.1 Experiments for Section 3

**Table 2: Hyperparameters for each workload.**
| Workload | Epochs/steps | Learning rate / decay | Weight decay | Batch size | $\rho$ search | Baseline optimizer |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Synthetic | 100 epochs | $0.1$ / step | $0.0$ | $128$ | $\{0.001,0.01,0.05,0.07,0.1,0.2,0.3,0.5,0.7,1.0,2.0\}$ | SGD |  |
| MNIST/MLP | $100$ epochs | $0.1$ / step | $0.0001$ | $128$ | $\{0.01,0.02,0.05,0.1,0.2\}$ | SGD with momentum $0.9$ |  |
| CIFAR-10/ResNet-18 | $200$ epochs | $0.1$ / step | $0.0005$ | $128$ | $\{0.001,0.005,0.01,0.02,0.05,0.1,0.2,0.5,1.0\}$ | SGD with momentum $0.9$ |  |
| ImageNet/ResNet-50 | $90$ epochs | $0.1$ / cosine | $0.0001$ | $512$ | $\{0.01,0.02,0.05,0.1,0.2\}$ | SGD with momentum $0.9$ |  |
| PoS tagging | $75000$ steps | $0.05$ / inverse square root | $0.1$ | $64$ | $\{0.01,0.02,0.05,0.1,0.2,0.3,0.5\}$ | AdamW ($\beta_{1}=0.9,\beta_{2}=0.98$) |  |
| Sentiment classification | $30$ epochs | $0.1$ / constant | $3e\text{-}6$ | $64$ | $\{0.01,0.02,0.05,0.1,0.2,0.3,0.5\}$ | SGD with momentum $0.8$ |  |
| Graph property prediction | $10^{5}$ steps | $0.001$ / constant | $0.0$ | $256$ | $\{0.01,0.02,0.05,0.1,0.2\}$ | Adam ($\beta_{1}=0.9,\beta_{2}=0.999$) |  |
| Atari game | $10^{7}$ steps | $2.5e\text{-}4$ / linear | $0.0$ | $256$ | $\{0.01,0.02,0.05,0.1,0.2\}$ | Adam ($\beta_{1}=0.9,\beta_{2}=0.999$) |  |
| CIFAR-10/ViT | $200$ epochs | $0.1$ / cosine | $0.0001$ | $128$ | $\{0.01,0.02,0.05,0.1,0.2\}$ | SGD with momentum $0.9$ |  |

**Table 3: Model scaling factors and values for each workload.**
| Workload | Scaling factor | Values |
| --- | --- | --- |
| Synthetic | # of neurons | $\{k*100|1\leq k\leq 10\}$ |
| MNIST/MLP | # of neurons | $\{[300*p,100*p]|p\in\{0.25,0.5,1,4,10\}\}$ |
| CIFAR-10/ResNet-18 | # of convolutional filters | $\{2^{k}|2\leq k\leq 8\}$ |
| ImageNet/ResNet-50 | # of convolutional filters | $\{16*k|1\leq k\leq 5\}$ |
| PoS tagging | dimension of hidden states | $\{128*k|1\leq k\leq 5\}$ |
| Sentiment classification | dimension of hidden states | $\{2^{k}|5\leq k\leq 9\}$ |
| Graph property prediction | # of neurons | $\{2^{k}|7\leq k\leq 9\}$ |
| Atari game | # of convolutional filters | $\{16*k|1\leq k\leq 4\}$ |
| CIFAR-10/ViT | dimension of hidden states | $\{2^{k}|5\leq k\leq 10\}$ |

For all the experiments in Section 3, we run the experiments with the same configurations over three different random seeds.
Most of the experiments are conducted with a single RTX3090 GPU having $24$GB memory while some experiments requiring larger memory are conducted with multiple RTX3090 GPUs.
The code is implemented with JAX (Bradbury et al., ) and provided as a supplementary material.
Many of our experiments and the hyperparameter values are based on examples provided by Flax (Heek et al., ) official repository.(^3^33[https://github.com/google/flax/tree/main/examples](https://github.com/google/flax/tree/main/examples))
The hyperparameter values and how the models are scaled for each workload are summarized in Table 2 and Table 3 respectively.
We present the additional details below.

#### Synthetic Regression / 2-layer MLP

We follow the student-teacher setting from Advani et al. () where the teacher is a randomly initialized $2$-layer ReLU network with $200$ neurons and the student is a $2$-layer ReLU network with a different number of neurons.
Each element for the input $x\in\mathbb{R}^{100}$ is sampled from a standard normal distribution while the target $y\in\mathbb{R}$ is calculated as the output of the teacher network added by Gaussian noise sampled from a standard normal distribution.
The models are trained on $20400$ training data, which is roughly the same as the number of parameters in the teacher model, and tested on the $5100$ data, which is a quarter of the number of the training data.

#### MNIST / 3-layer MLP

We train LeNet-300-100 (LeCun et al., ) for the MNIST (LeCun et al., ).
The learning rate decays by $0.1$ after $50\%$ and $75\%$ of the total epochs.
We scale the models while preserving the relative proportions of the number of neurons in each layer as $3:1$.

#### CIFAR-10 / ResNet-18

We train ResNet-18 (He et al., ) for the CIFAR-10 (Krizhevsky et al., ).
We choose the hyperparameters as similar to Andriushchenko and Flammarion ().
The learning rate decays by $0.1$ after $50\%$ and $75\%$ of the total epochs.

#### ImageNet / ResNet-50

We train ResNet-50 (He et al., ) for the ImageNet (Deng et al., ).
We choose the hyperparameters as similar to Du et al. () and use a linear warmup of $5000$ steps.
We additionally experiment with $\rho=0.005$ for the two smallest models.

#### PoS tagging/ Transformer

We train Encoder-only Transformer (Vaswani et al., ) for the Universal Dependencies (Nivre et al., ) – Ancient Greek.
We use a linear warmup of $8000$ steps.
We evaluate the validation accuracy once every $1000$ step and report the best value except for the experiment in Figure 6(d).
The dimension of MLP and the number of attention heads are scaled as $4\times$ and $1/64\times$ of the dimension of the hidden states following the Flax example.

#### SST / LSTM

We train LSTM (Hochreiter and Schmidhuber, ) for SST2 (Socher et al., ) where the task is a binary classification (positive/negative) of the movie reviews.
We evaluate the validation accuracy for every epoch and report the best value.
The embedding size is scaled as $300/256\times$ of the dimension of hidden states following the Flax example.

#### Graph property prediction / GCN

We train $2$-layer Graph Convolutional Networks (Kipf and Welling, ) for the ogbg-molpcba (Hu et al., ).
Here, the input is a graph of a molecule where nodes and edges each represent atoms and chemical bonds.
The task is a binary classification of whether a molecule inhibits HIV replication or not.

#### Atari game / CNN

We train $5$-layer CNNs for the Atari Breakout-v5 game (Mnih et al., ).
We train the Actor-Critic networks (Konda and Tsitsiklis, ) with proximal policy optimization (Schulman et al., ).
We evaluate the validation score once every $100$ step and report the best value.
We also use gradient clipping of $0.5$ for all models.

#### CIFAR-10 / ViT

For the experiment in Figure 6(e), we train $6$-layer Vision Transformers (Dosovitskiy et al., ) for the CIFAR-10 (Krizhevsky et al., ) using the patch size of $4\times 4$.
We scale the dimension of MLP and the number of attention heads as $2\times$ and $1/32\times$ of the dimension of hidden states.

### A.2 Experiments for Section 4

For the experiments in Figures 2 and 3(a), we follow the setting in Andriushchenko and Flammarion ().(^4^44[https://github.com/tml-epfl/understanding-sam/tree/main/one_layer_relu_nets](https://github.com/tml-epfl/understanding-sam/tree/main/one_layer_relu_nets))
Specifically, we train one-hidden-layer ReLU networks where each data has input $x\in\mathbb{R}$ and target $y\in\mathbb{R}$.
Here, the networks are trained on the quadratic loss with full-batch GD or SAM with $\rho=0.3$.
Additionally, the optimization trajectories in Figure 3(a) are plotted following Li et al. ().(^5^55[https://github.com/tomgoldstein/loss-landscape](https://github.com/tomgoldstein/loss-landscape))
Specifically, the trajectories are plotted along the PCA directions calculated from converged minima of two different paths from SGD and one path from SAM.

### A.3 Experiments for Section 5

#### Linear stability

For the experiments of Figures 5(a), 3(b) and 3(c), we follow the setting in Wu et al. ().
Specifically, we set up $3$-layer MLP having $[3000,1000]$ hidden neurons with squared loss, so that the local quadratic approximation becomes precise, and train the networks on MNIST.
We use $1000$ random samples to calculate the non-uniformity, and all models are trained to reach near zero loss.
The networks are trained with a constant learning rate of $0.1$ without weight decay or momentum.

#### Convergence – Matrix Factorization

For the matrix factorization experiment in Figure 5(b), we solve the following non-convex regression problem: $\min_{W_{1},W_{2}}\mathbb{E}_{x\sim\mathcal{N}(0,I)}{\|W_{2}W_{1}x-Ax\|^{2}}$ where the objective function is smooth and satisfies the PL-condition (Loizou et al., ).
We choose $A\in\mathbb{R}^{10\times 6}$ and generate 1000 training samples, which are used for training a rank $k$ linear network with two matrix factors $W_{1}\in\mathbb{R}^{k\times 6}$ and $W_{2}\in\mathbb{R}^{10\times k}$. Here, interpolation is satisfied when rank $k=10$.
We train two linear networks with $k\in\{4,10\}$ for $100$ epochs with a constant learning rate of $0.0005$ and compare the convergence speed.

#### Generalization

For the experiment of Figure 5(c), we follow the setup of Allen-Zhu et al. ().
Specifically, we train $2$-layer ReLU networks with synthetic data.
Here, each element of the input $x=(x_{1},x_{2},x_{3},x_{4})\in\mathbb{R}^{4}$ for synthetic data is sampled from random Gaussian distribution and then normalized to satisfy $\|x\|_{2}=1$, and target $y$ is calculated as $y=(\sin(3x_{1})+\sin(3x_{2})+\sin(3x_{3})-2)^{2}\cdot\cos(7x_{4})$.
The weights and biases of the first layer are initialized from $\mathcal{N}(0,1/m)$ where $m$ is the number of hidden neurons, and the weights of the second layer are initialized from $\mathcal{N}(0,1)$.
We only train the weights of the first layer for $800$ epochs, while the biases of the first layer and the weights of the second layer are frozen to initialized values.
We use $1000$ and $5000$ data points for training and testing respectively.
We use a batch size of $50$ without weight decay and decay learning rate by $0.1$ after $50\%$ of the total epochs.
We perform the grid search over learning rate and $\rho$ from $\{10^{-k}|2\leq k\leq 7\}$ and $\{10^{-k}|1\leq k\leq 5\}$ respectively.

## Appendix B Additional experimental results

In this section, we present additional experimental results.

#### More results on overparameterization benefit of SAM

Figure: (a) Improvement in validation metrics by SAM
Refer to caption: x32.png

We present the full experimental results of the effect of overparameterization on SAM including the absolute metrics in Figure 7.
There is a consistent trend that SAM improves with overparameterization in all tested cases.

#### More results on optimal perturbation bound

Figure: Figure 8: Validation accuracy versus $\rho$ for MNIST and $3$-layer MLP.
Refer to caption: x48.png

Figure: Figure 9: Validation accuracy versus $\rho$ for ResNet-18 and CIFAR-10.
Refer to caption: x17.png

Figure: Figure 10: Validation accuracy versus $\rho$ for ResNet-50 and ImageNet.
Refer to caption: x57.png

Figure: Figure 11: Validation accuracy versus $\rho$ for LSTM and SST2.
Refer to caption: x63.png

We plot the validation accuracy of SAM versus different values of $\rho$ in Figures 8, 9, 10 and 11.
It is observed that $\rho^{\star}$ tends to increase as the model becomes more overparameterized;
on CIFAR-10 with ResNet18, the smallest model has $\rho^{\star}=0.01$ while the largest three have $\rho^{\star}=0.2$.

#### More results on label noise

Figure: (a) Effect of label noise
Refer to caption: x70.png

More results on the effect of overparameterization on SAM under label noise are presented in Figure 12.
Overall, we find SAM benefits from overparameterization significantly more than SGD in the presence of label noise.
Precisely, the accuracy improvement made by SAM keeps on increasing as the model becomes more overparameterized, and this trend is more pronounced with higher noise levels;
e.g., it rises from $5\%$ to nearly $50\%$ at the highest noise rate.

#### More results on sparse overparameterization

Figure: (a) Random pruning
Refer to caption: x74.png

Figure: (a) Random pruning
Refer to caption: x78.png

Figure: (a) MNIST/MLP
Refer to caption: x82.png

Additional results on the effect of sparsification on the generalization benefit of SAM are plotted in Figures 13 and 14.
Here, we try two sparsification methods that do not require pertaining, random pruning and SNIP (Lee et al., ).
For both methods, we note that the generalization improvement by SAM tends to increase as the model becomes more sparsely overparameterized.

We also plot the effect of sparsification on $\rho^{\star}$ in Figure 15.
We find that $\rho^{\star}$ is sometimes different between small dense and large sparse models despite having a similar number of parameters;
for the MLP of $61$k parameters on MNIST, $\rho^{\star}$ changes over different sparsity levels and sparsification methods, but this does not generalize to the CIFAR-10 and ResNet-18.
This indicates that it is not just the parameter count that affects the behavior of SAM, but some other factors such as the pattern of parameterization also have an influence on how SAM shapes training.

#### More results on regularization

Figure: (a) w/o weight decay
Refer to caption: x86.png

More results on the effect of regularization on SAM are presented in Figure 16.
We find that overparameterization does not increase the generalization benefit of SAM.
We suspect this is because the models are prone to overfitting in these cases and overparameterizing models may decrease the overall performance both for SGD and SAM;
for example in Figure 16(c), the validation accuracy drops after $11.2$m parameters.

## Appendix C Results on SAM under Linearized Regime

Figure: Figure 17: Effect of linearization on SAM. Here, $\alpha$ controls the degree of linearization. High linearization does not yield improvement, and in fact, SAM ($\rho=0.001$) underperforms SGD in the linearized regime (left), although both achieve effective linearization at $\alpha=1000$ with stability close to $1$ (right). See Appendix C for full details.
Refer to caption: x92.png

Recent studies suggest that highly overparameterized models can behave like linearized networks (Jacot et al., ), while such implicit linearization phenomenon can coincide independently of overparameterization (Chizat et al., ).
One might wonder if the increased effectiveness of SAM directly comes from the overparameterization itself or is rather due to linearization.
To verify, we reproduce experiments in Chizat et al. () and see how SAM performs in the linearized regimes while fixing the number of parameters.
Specifically, we train VGG-11 (Simonyan and Zisserman, ) on the Cifar-10 with the $\alpha$-scaled squared loss $L(x,y)=\|f(x)-y/\alpha\|^{2}$ and use the centered model whose initial output is set to 0 0.
Here, a large value of $\alpha$ leads to a higher degree of linearization of the models.

The results are reported in Figure 17.
We observe that SAM underperforms SGD in the linearized regimes;
while SAM ($\rho=0.001$) and SGD both achieve effective linearization at $\alpha=1000$, SAM underperforms SGD by more than $10\%$.
This indicates that linearization is not the main factor, and overparameterization itself is what leads to the improvement of SAM in previous experiments.

## Appendix D Empirical Measurement of Lipschitz Smoothness and PL Constants

We measure the empirical Lipschitz smoothness constant and PL constant based on Zhang et al. ().
Specifically, the empirical smoothness $\hat{\beta}(x_{k})$ and empirical PL constant $\hat{\alpha}(x_{k})$ at iteration $k$ is computed as follows:

$$ $\displaystyle\hat{\beta}(x_{k})$ $\displaystyle=\max_{\gamma\in\{\delta,2\delta,\dots,1\}}\frac{\|\nabla f(x_{k} +\gamma\mathbf{d})-\nabla f(x_{k})\|_{2}}{\|\gamma\mathbf{d}\|_{2}},$ (7) $\displaystyle\hat{\alpha}(x_{k})$ $\displaystyle=\min_{\gamma\in\{\delta,2\delta,\dots,1\}}\frac{\|\nabla f(x_{k} +\gamma\mathbf{d})\|_{2}^{2}}{f(x_{k}+\gamma\mathbf{d})-f(x^{\star})}.$ (8) $$

where $\mathbf{d}=x_{k+1}-x_{k}$ and $\delta\in(0,1)$ where we choose $\delta=0.1$.
We measure these quantities at the end of every epoch throughout training.
The results are shown in Figure 18.

Figure: (a) Lipschitz smoothness
Refer to caption: x94.png

## Appendix E Proof of Theorem 5.6

In this section, we show that a stochastic SAM converges linearly under an overparameterized regime.
To put into perspective, this is the rate of convergence of gradient descent for a family of functions satisfying the PL-condition and smoothness assumptions (Karimi et al., ).
We first make several remarks on this result below.

- •
Crucially, this result shows that with overparameterization, a stochastic SAM can converge as fast as the deterministic gradient method at a linear convergence rate.
It is much faster than the well-known sublinear rate of $\mathcal{O}(1/t)$ for SAM (Andriushchenko and Flammarion, ).
- •
When $\rho=0$, we recover the well-known convergence rate for SGD in interpolated regime (Bassily et al., ).
- •
This result does not require the bounded variance assumption (Andriushchenko and Flammarion, ) since the interpolation provides necessary guarantees.
This suggests that overparameterization can ease the convergence of SAM.

We prove the convergence for an unnormalized mini-batch SAM given as

$$ $x_{t+1}=x_{t}-\eta g_{t}^{{}_{B}}(x_{t}+\rho g_{t}^{{}_{B}}(x_{t})),$ $$

where $g_{t}^{{}_{B}}(x)=\frac{1}{B}\sum_{i\in I_{t}^{B}}\nabla f_{i}(x)$ and $I_{t}^{B}\subseteq\{1,...,n\}$ is a set of indices for data points in the mini-batch of size $B$ sampled at step $t$.
This is a more general stochastic variant of SAM where a stochastic SAM in Section 5.2 is a particular case of a mini-batch SAM with mini-batch size $B=1$.

We first make the following assumptions:

($\beta$-smothness of $f_{i}$). There exists $\beta>0$ s.t. $\|\nabla f_{i}(x)-\nabla f_{i}(y)\|\leq\beta\|x-y\|$ for all $x,y\in\mathbb{R}^{d}$,

($\lambda$-smothness of $f$). There exists $\lambda\!>\!0$ such that $\|\nabla f(x)-\nabla f(y)\|\leq\lambda\|x-y\|$ for all $x,y\in\mathbb{R}^{d}$,

($\alpha$-PLness of $f$). There exists $\alpha>0$ s.t. $\|\nabla f(x)\|^{2}\geq\alpha(f(x)-f(x^{\star}))$ for all $w,v\in\mathbb{R}^{d}$,

(Interpolation). If $f(x^{\star})=0$ and $~{}\nabla f(x^{\star})=0$, then $f_{i}(x^{\star})=0$ and $\nabla f_{i}(x^{\star})=0$ for $i=1,\ldots,n$, where $n$ is the number of training data points.

Before we prove the main theorem, we first introduce two lemma important to the proof.

###### Lemma E.1 .

Suppose that Assumption (A1) holds.
Then

$$ $\displaystyle\langle\nabla f_{i}(x_{t+1/2}),\nabla f(x_{t})\rangle$ $\displaystyle\geq\langle\nabla f_{i}(x_{t}),\nabla f(x_{t})\rangle-\frac{{ \beta}\rho}{2}\|\nabla f_{i}(x_{t})\|^{2}-\frac{\beta\rho}{2}\|\nabla f(x_{t}) \|^{2},$ (9) $$

where $x_{t+1/2}=x_{t}+\rho\nabla f_{i}(x_{t})$.

This lemma shows how well a stochastic SAM gradient $\nabla f_{i}(x_{t+1/2})$ aligns with the true gradient $\nabla f(x_{t})$.
The two gradients become less aligned as $\beta$ and $\rho$ grow bigger, *i.e*. for sharper landscape and larger perturbation size.

###### Proof.

We first add and subtract $\nabla f_{i}(x_{t})$ on the left side of the inner product

$$ $\displaystyle\langle\nabla f_{i}(x_{t+1/2}),\nabla f(x_{t})\rangle=\underbrace {\langle\nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t}),\nabla f(x_{t})\rangle}_{ \tau_{1}}+\langle\nabla f_{i}(x_{t}),\nabla f(x_{t})\rangle.$ (10) $$

We here bound the term $\tau_{1}$ so that it becomes an equality when $\rho=0$.
To achieve this, we start with the following binomial square, which is trivially lower bounded by 0.

$$ $\displaystyle 0$ $\displaystyle\leq\frac{1}{2}\|\nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t})+ \beta\rho\nabla f(x_{t})\|^{2}$ $$

We then expand the above binomial square so that the term containing $\tau_{1}$ appears.

$$ $\displaystyle 0$ $\displaystyle\leq\frac{1}{2}\|\nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t})\|^{2 }+\underbrace{\langle\nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t})~{},~{}\beta \rho\nabla f(x_{t})\rangle}_{\beta\rho\tau_{1}}\,+\,\frac{1}{2}\|\beta\rho \nabla f(x_{t})\|^{2}$ $$

We subtract the term $\beta\rho\tau_{1}$ on both sides of the inequality which gives

$$ $\displaystyle-\langle\nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t})~{},~{}\beta \rho\nabla f(x_{t})\rangle$ $\displaystyle\leq\frac{1}{2}\|\nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t})\|^{2 }+\frac{\beta^{2}\rho^{2}}{2}\|\nabla f(x_{t})\|^{2}.$ $$

Then we upper bound the right-hand side using the Assumption (A1):

$$ $\displaystyle-\langle\nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t})~{},~{}\beta \rho\nabla f(x_{t})\rangle$ $\displaystyle\leq\frac{{\beta}^{2}}{2}\|x_{t+1/2}-x\|^{2}+\frac{\beta^{2}\rho^ {2}}{2}\|\nabla f(x_{t})\|^{2}$ $\displaystyle=\frac{{\beta}^{2}\rho^{2}}{2}\|\nabla f_{i}(x_{t})\|^{2}+\frac{ \beta^{2}\rho^{2}}{2}\|\nabla f(x_{t})\|^{2}.$ $$

We divide both sides with $\beta\rho$, obtaining:

$$ $\displaystyle-\langle\nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t}),\nabla f(x_{t })\rangle$ $\displaystyle\leq\frac{{\beta}\rho}{2}\|\nabla f_{i}(x_{t})\|^{2}+\frac{\beta \rho}{2}\|\nabla f(x_{t})\|^{2}.$ $$

Applying this to (10) gives the bound in the lemma statement.
∎

###### Lemma E.2 .

Suppose that Assumption (A1) holds.
Then

$$ $\left\|\nabla f_{i}(x_{t+{1/2}})\right\|^{2}\leq(\beta\rho+1)^{2}\|\nabla f_{i }(x_{t})\|^{2},$ (11) $$

where $x_{t+1/2}=x_{t}+\rho\nabla f_{i}(x_{t})$.

This second lemma shows that the norm of a stochastic SAM gradient is bounded by the norm of the stochastic gradient.
Similar to the Lemma E.1, as $\beta$ and $\rho$ grow bigger the norm for a stochastic SAM gradient can become larger than the norm of the true gradient.

###### Proof.

We use the following binomial squares:

$$ $\begin{split}&\|\nabla f_{i}(x_{t+1/2})\|^{2}\\ &=\|\nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t})\|^{2}+2\langle\nabla f_{i}(x_{ t+1/2})-\nabla f_{i}(x_{t}),\nabla f_{i}(x_{t})\rangle+\|\nabla f_{i}(x_{t})\| ^{2}.\end{split}$ $$

We bound the right-hand side using Cauchy-Schwarz inequality and Assumption (A1), which gives

$$ $\displaystyle\left\|\nabla f_{i}(x_{t+{1/2}})\right\|^{2}$ $\displaystyle=\|\nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t})\|^{2}+2\langle \nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t}),\nabla f_{i}(x_{t})\rangle+\| \nabla f_{i}(x_{t})\|^{2}$ $\displaystyle\underset{\text{C.S.}}{\leq}\|\nabla f_{i}(x_{t+1/2})-\nabla f_{i }(x_{t})\|^{2}+2\|\nabla f_{i}(x_{t+1/2})-\nabla f_{i}(x_{t})\|\|\nabla f_{i}( x_{t})\|+\|\nabla f_{i}(x_{t})\|^{2}$ $\displaystyle\underset{\hyperlink{assump:betasmoo}{\textbf{(A1)}}}{\leq}~{} \beta^{2}\|x_{t+1/2}-x_{t}\|^{2}+2\beta\|x_{t+1/2}-x_{t}\|\|\nabla f_{i}(x_{t} )\|+\|\nabla f_{i}(x_{t})\|^{2}~{}$ $\displaystyle~{}=~{}~{}\beta^{2}\rho^{2}\|\nabla f_{i}(x_{t})\|^{2}+2\beta\rho \|\nabla f_{i}(x_{t})\|^{2}+\|\nabla f_{i}(x_{t})\|^{2}~{}$ $\displaystyle~{}=~{}(\beta\rho+1)^{2}\|\nabla f_{i}(x_{t})\|^{2}$ $$

∎

These two lemmas essentially show how similar a stochastic SAM gradient is to the stochastic gradient, where the two become more similar as $\beta$ and $\rho$ decrease, which aligns well with our intuition.
Using Lemma E.1 and E.2, we provide the convergence result in the following theorem.

###### Theorem E.3 .

Suppose that Assumptions (A1-4) holds.
For any mini-batch size $B\in\mathbb{N}$ and $\rho\leq\frac{1}{(\beta/\alpha+1/2)\beta}$, unnormalized mini-batch SAM with constant step size $\eta^{\star}_{B}\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{1-(\kappa_{B}+1/
2)\beta\rho}{2\lambda\kappa_{B}(\beta\rho+1)^{2}}$ gives the following guarantee at step $t$:

$$ $\underset{x_{t}}{\mathbb{E}}\left[f(x_{t})\right]\leq\left(1-\frac{\alpha\, \eta^{\star}_{B}}{2}\Big{(}1-\Big{(}\kappa_{B}+\frac{1}{2}\Big{)}\beta\rho\Big {)}\right)^{t}\,f(x_{0}),$ $$

where $\kappa_{B}=\frac{1}{B}\left(\frac{B-1}{2}+\frac{\beta}{\alpha}\right)$.

This theorem states that mini-batch SAM converges at a linear rate under overparameterization.

###### Proof.

Proof can be outlined in 3 steps.

Handle terms containing mini-batch SAM gradient $g_{t}^{{}_{B}}(x_{t}+\rho g_{t}^{{}_{B}}(x_{t}))$ using bounds from (A1).

Take conditional expectation $\mathbb{E}{\left[\,\cdot\,|x_{t}\right]}$ and substitute expectation of function of mini-batch gradient $g_{t}^{{}_{B}}$ with terms containing $\|\nabla f(x_{t})\|$ and $\mathbb{E}{\left[\|\nabla f_{i}(x_{t})\|^{2}~{}\Big{\lvert}~{}x_{t}\right]}$.

Bound the two terms from step 2, one using Assumptions (A1) and (A4) and the other using Assumption (A3) and (A4) which results in all the terms to contain $f(x_{t})$.
Then finally we take total expectations to derive the final runtime bound.

We start from the quadratic upper bound derived from Assumption (A2);

$$ $f(x_{t+1})\leq f(x_{t})+\langle\nabla f(x_{t}),~{}x_{t+1}-x_{t}\rangle+\frac{ \lambda}{2}\|x_{t+1}-x_{t}\|^{2}.$ $$

Applying mini-batch SAM update, we then have

$$ $f(x_{t})-f(x_{t+1})\geq\eta\left\langle~{}\nabla f(x_{t})~{},~{}g_{t}^{{}_{B}} (x_{t+{1/2}})~{}\right\rangle-\frac{\eta^{2}\lambda}{2}\left\|g_{t}^{{}_{B}}(x _{t+{1/2}})\right\|^{2},$ $$

where $x_{t+1/2}=x_{t}+\rho g_{t}^{{}_{B}}(x_{t})$.

step 1.   We can see that there are two terms that contain a mini-batch SAM gradient $g_{t}^{{}_{B}}(x_{t+1/2})$.
We see that each can be bounded directly using Lemma E.1 and E.2, which gives

$$ $\displaystyle f(x_{t})-f(x_{t+1})$ $\displaystyle\geq\eta\left(\langle g_{t}^{{}_{B}}(x_{t}),\nabla f(x_{t}) \rangle-\frac{{\beta}\rho}{2}\|g_{t}^{{}_{B}}(x_{t})\|^{2}-\frac{\beta\rho}{2} \|\nabla f(x_{t})\|^{2}\right)$ $\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}-\frac{\eta^{2}\lambda}{2}~{}(\beta \rho+1)^{2}~{}\|g_{t}^{{}_{B}}(x_{t})\|^{2}$ $\displaystyle~{}~{}~{}=~{}~{}\eta\langle g_{t}^{{}_{B}}(x_{t}),\nabla f(x_{t}) \rangle-\frac{\eta\beta\rho}{2}\|\nabla f(x_{t})\|^{2}-\frac{\eta}{2}\left( \eta\lambda~{}(\beta\rho+1)^{2}+\beta\rho\right)~{}\|g_{t}^{{}_{B}}(x_{t})\|^{ 2}.$ $$

step 2.   Now we apply $\mathbb{E}{\left[\,\cdot\,|x_{t}\right]}$ to all the terms.

$$ $\displaystyle\mathbb{E}{\left[f(x_{t})-f(x_{t+1})~{}\big{\lvert}~{}x_{t}\right]}$ $\displaystyle=f(x_{t})-\mathbb{E}{\left[f(x_{t+1})~{}\big{\lvert}~{}x_{t} \right]}$ $\displaystyle\geq\eta\mathbb{E}{\left[\langle g_{t}^{{}_{B}}(x_{t}),\nabla f(x _{t})\rangle~{}\Big{\lvert}~{}x_{t}\right]}-\frac{\eta\beta\rho}{2}\mathbb{E}{ \left[\|\nabla f(x_{t})\|^{2}~{}\big{\lvert}~{}x_{t}\right]}$ $\displaystyle\ \ \ \ \ \ -\frac{\eta}{2}\left(\eta\lambda~{}(\beta\rho+1)^{2}+ \beta\rho\right)~{}\mathbb{E}{\left[\|g_{t}^{{}_{B}}(x_{t})\|^{2}~{}\Big{ \lvert}~{}x_{t}\right]}$ $\displaystyle=\eta\left(1-\frac{\beta\rho}{2}\right)\|\nabla f(x_{t})\|^{2}- \frac{\eta}{2}\left(\eta\lambda~{}(\beta\rho+1)^{2}+\beta\rho\right)~{}\mathbb {E}{\left[\|g_{t}^{{}_{B}}(x_{t})\|^{2}~{}\Big{\lvert}~{}x_{t}\right]}.$ $$

Here we expand the term $\mathbb{E}{\left[\|g_{t}^{{}_{B}}(x_{t})\|^{2}~{}\Big{\lvert}~{}x_{t}\right]}$ by expanding the mini-batched function into individual function estimators as follows.

$$ $\begin{split}&\mathbb{E}_{g^{{}_{B}}_{t}}{\left[\left\|g_{t}^{{}_{B}}(x_{t}) \right\|^{2}~{}\Big{\lvert}~{}x_{t}\right]}\\ &=\mathbb{E}_{I_{t}^{B}}{\left[\left\langle\frac{1}{B}\sum_{i\in I_{t}^{B}} \nabla f_{i}(x_{t})~{},~{}\frac{1}{B}\sum_{j\in I_{t}^{B}}\nabla f_{j}(x_{t}) \right\rangle~{}\Bigg{\lvert}~{}x_{t}\right]}\\ &=\frac{1}{B^{2}}\left\{\sum_{i\in I_{t}^{B}}\mathbb{E}_{f_{i}}{\left[\left\| \nabla f_{i}(x_{t})\right\|^{2}~{}\Big{\lvert}~{}x_{t}\right]}+\sum_{i\in I_{t }^{B}}\sum_{\begin{subarray}{c}j\in I_{t}^{B}\\ \text{(}j\neq i\text{)}\end{subarray}}\mathbb{E}_{f_{i},f_{j}}{\left[\langle \nabla f_{i}(x_{t}),\nabla f_{j}(x_{t})\rangle~{}\Big{\lvert}~{}x_{t}\right]} \right\}\\ &=\frac{1}{B}\mathbb{E}{\left[\left\|\nabla f_{i}(x_{t})\right\|^{2}~{}\Big{ \lvert}~{}x_{t}\right]}+\frac{B-1}{B}\left\|\nabla f(x_{t})\right\|^{2}.\end{split}$ (12) $$

Using (12), we get

$$ $\displaystyle f(x_{t})-\mathbb{E}{\left[f(x_{t+1})~{}\big{\lvert}~{}x_{t} \right]}$ $\displaystyle\geq\eta\left(1-\frac{\beta\rho}{2}\right)\|\nabla f(x_{t})\|^{2}$ $\displaystyle~{}~{}~{}-\frac{\eta}{2}\left(\eta\lambda~{}(\beta\rho+1)^{2}+ \beta\rho\right)\left(\frac{1}{B}\mathbb{E}{\left[\left\|\nabla f_{i}(x_{t}) \right\|^{2}~{}\Big{\lvert}~{}x_{t}\right]}+\frac{B-1}{B}\left\|\nabla f(x_{t} )\right\|^{2}\right)$ $\displaystyle=\eta\left(\left(1-\frac{\beta\rho}{2}\right)-\frac{B-1}{2B}\left (\eta\lambda(\beta\rho+1)^{2}+\beta\rho\right)\right)\|\nabla f(x_{t})\|^{2}$ $\displaystyle~{}~{}~{}-\frac{\eta}{2B}\left(\eta\lambda(\beta\rho+1)^{2}+\beta \rho\right)\mathbb{E}{\left[\left\|\nabla f_{i}(x_{t})\right\|^{2}~{}\Big{ \lvert}~{}x_{t}\right]}.$ (13) $$

step 3.   In this step, we bound the two terms and take the total expectation to derive the final runtime bound.

We first derive a bound for $\mathbb{E}{\left[\left\|\nabla f_{i}(x_{t})\right\|^{2}~{}\Big{\lvert}~{}x_{t}
\right]}$.
We start from the following bound derived from Assumption (A1):

$$ $\|\nabla f_{i}(x_{t})-\nabla f_{i}(x^{\star})\|^{2}\leq 2\beta(f_{i}(x_{t})-f_ {i}(x^{\star})).$ $$

By Assumption (A4), this reduces to

$$ $\|\nabla f_{i}(x_{t})\|^{2}\leq 2\beta f_{i}(x_{t}).$ $$

Applying this to (13) gives

$$ $\displaystyle f(x_{t})-\mathbb{E}{\left[f(x_{t+1})~{}\big{\lvert}~{}x_{t} \right]}$ $\displaystyle\geq\eta\left(\left(1-\frac{\beta\rho}{2}\right)-\frac{B-1}{2B} \left(\eta\lambda(\beta\rho+1)^{2}+\beta\rho\right)\right)\|\nabla f(x_{t})\|^ {2}$ $\displaystyle~{}~{}~{}~{}~{}-\frac{\eta\beta}{B}\left(\eta\lambda(\beta\rho+1) ^{2}+\beta\rho\right)\mathbb{E}{\left[f_{i}(x_{t})\lvert x_{t}\right]}$ $\displaystyle=\eta\underbrace{\left(\left(1-\frac{\beta\rho}{2}\right)-\frac{B -1}{2B}\left(\eta\lambda(\beta\rho+1)^{2}+\beta\rho\right)\right)}_{\tau_{2}} \|\nabla f(x_{t})\|^{2}$ $\displaystyle~{}~{}~{}~{}~{}-\frac{\eta\beta}{B}\left(\eta\lambda(\beta\rho+1) ^{2}+\beta\rho\right)f(x_{t}).$ (14) $$

Next, to bound $\|\nabla f(x_{t})\|^{2}$, we use the following bound derived from applying $f(x^{*})=0$ from (A4) to (A3):

$$ $\|\nabla f(x)\|^{2}\geq\alpha f(x).$ (15) $$

Assuming $\tau_{2}\geq 0$ which we provide a sufficient condition at the end of the proof, we apply (15) to (14) which gives

$$ $\displaystyle f(x_{t})-\mathbb{E}{\left[f(x_{t+1})~{}\big{\lvert}~{}x_{t} \right]}$ $\displaystyle\geq\eta\alpha\left(\left(1-\frac{\beta\rho}{2}\right)-\frac{B-1} {2B}\left(\eta\lambda(\beta\rho+1)^{2}+\beta\rho\right)\right)f(x_{t})-\frac{ \eta\beta}{B}\left(\eta\lambda(\beta\rho+1)^{2}+\beta\rho\right)f(x_{t})$ $\displaystyle=\eta\left(\alpha-\alpha\bigg{(}\underbrace{\frac{1}{B}\Big{(} \frac{B-1}{2}+\frac{\beta}{\alpha}\Big{)}}_{\kappa_{B}}+\frac{1}{2}\bigg{)} \beta\rho-\eta(\beta\rho+1)^{2}\underbrace{\frac{\lambda}{B}\Big{(}\alpha\frac {B-1}{2}+\beta\Big{)}}_{\lambda\alpha\kappa_{B}}\right)f(x_{t})$ $\displaystyle=\eta\alpha\left(1-\left(\kappa_{B}+\frac{1}{2}\right)\beta\rho- \eta\lambda(\beta\rho+1)^{2}\kappa_{B}\right)f(x_{t}).$ $$

Hence, we get

$$ $\displaystyle\mathbb{E}{\left[f(x_{t+1})~{}\big{\lvert}~{}x_{t}\right]}$ $\displaystyle\leq\left(1-\eta\alpha\Big{(}1-\Big{(}\kappa_{B}+\frac{1}{2}\Big{ )}\beta\rho\Big{)}+\eta^{2}\alpha\lambda(\beta\rho+1)^{2}\kappa_{B}\right)f(x_ {t}).$ $$

Applying total expectation on both sides gives

$$ $\displaystyle\mathbb{E}{\left[f(x_{t+1})\right]}\leq\left(1-\eta\alpha\Big{(}1 -\Big{(}\kappa_{B}+\frac{1}{2}\Big{)}\beta\rho\Big{)}+\eta^{2}\alpha\lambda( \beta\rho+1)^{2}\kappa_{B}\right)\mathbb{E}{\left[f(x_{t})\right]}.$ (16) $$

Optimizing the multiplicative term in (16) with respect to $\eta$ gives $\eta=\frac{1-(\kappa_{B}+1/2)\beta\rho}{2\lambda\kappa_{B}(\beta\rho+1)^{2}}$, which is $\eta^{\star}_{B}$ in the theorem statement.
With assumption of $\rho\leq\frac{1}{(\beta/\alpha+1/2)\beta}$ so that we have $\eta^{\star}_{B}\geq 0$, applying this to (16) gives

$$ $\underset{}{\mathbb{E}}\left[f(x_{t+1})\right]\leq\left(1-\frac{\alpha\,\eta^{ \star}_{B}}{2}\Big{(}1-\Big{(}\kappa_{B}+\frac{1}{2}\Big{)}\beta\rho\Big{)} \right)\,\underset{}{\mathbb{E}}\left[f(x_{t})\right],$ $$

which provides the desired convergence rate.

Last but not least, we calculate the upper bound for $\rho$ to satisfy the assumption $\tau_{2}\geq 0$ by substituting $\eta$ for $\eta^{\star}_{B}$ in $\tau_{2}$, yielding $\rho\leq\frac{2B\kappa_{B}+2\beta/\alpha}{(2B-1)\kappa_{B}+\beta/\alpha}\frac{
1}{\beta}$.
Minimizing this upper bound with respect to $B$ gives $\rho\leq\frac{1}{\beta}$, which is a looser bound than $\rho\leq\frac{1}{(\beta/\alpha+1/2)\beta}$.
∎

###### Lemma E.1 .

###### Proof.

###### Lemma E.2 .

###### Proof.

###### Theorem E.3 .

###### Proof.

## Appendix F Proof of Theorem 5.7

In this section, we provide the formal version of Theorem 5.7 and its proof.

### F.1 Notation and Setup

Throughout this section, we use the same notations and setups as Allen-Zhu et al. ().
We remark that the notations are different from those used in Sections 5.3, E and G.

First, let us assume the unknown data distribution $\mathcal{D}$ where each data $z=(x,y)$ consists of the input $x\in\mathbb{R}^{d}$ and the corresponding label $y\in\mathcal{Y}$.
We also assume, without loss of generality, that $\|x\|_{2}=1$ and $x_{d}=1/2$.
The loss function $L:\mathbb{R}^{k}\times\mathcal{Y}\rightarrow\mathbb{R}$ is assumed to be non-negative, convex, $1$-Lipschitz continuous, and $1$-smooth with respect to its first argument.

Next, we define the target network $F^{*}=(f^{*}_{1},\cdots,f^{*}_{k}):\mathbb{R}^{d}\rightarrow\mathbb{R}^{k}$ as

$$ $f^{*}_{r}(x)\stackrel{{\scriptstyle\rm def}}{{=}}\sum_{i=1}^{p}a^{*}_{r,i}\phi _{i}(\langle w^{*}_{1,i},x\rangle)\langle w^{*}_{2,i},x\rangle$ (17) $$

where each $\phi_{i}:\mathbb{R}\rightarrow\mathbb{R}$ is an infinite-order smooth function.
Here, we assume that $\|w^{*}_{1,i}\|_{2}=\|w^{*}_{2,i}\|_{2}=1,|a^{*}_{r,i}|\leq 1$ hold for all $i\in\{1,\cdots,p\}$.
We denote the sample and network complexity of $\phi$ as $\mathfrak{C}_{\mathfrak{s}}$ and $\mathfrak{C}_{\epsilon}$ respectively (see Section 2 of Allen-Zhu et al. () for the formal definitions).
Suppose we have a concept class $\mathcal{C}$ that consists of all functions $F^{*}$ with bounded number of parameters $p$ and complexity $\mathfrak{C}$.
We also denote the population risk achieved by the best target function $F^{*}$ in this concept class as OPT, *i.e*., $\textsf{{OPT}}=\underset{F^{\star}\in\mathcal{C}}{\min}\mathbb{E}_{(x,y)\sim
\mathcal{D}}[L(F^{*}(x),y]$

Then, we define the learner network $F=(f_{1},\cdots,f_{k}):\mathbb{R}^{d}\rightarrow\mathbb{R}^{k}$ as below.

$$ $f_{r}(x)\stackrel{{\scriptstyle\rm def}}{{=}}\sum_{i=1}^{m}a^{(0)}_{r,i} \operatorname{ReLU}(\langle w_{i},x\rangle+b^{(0)}_{i}).$ (18) $$

Note that the learner network is a $2$-layer ReLU network with $m$ neurons.
We train the network with $n$ sampled data sampled from $\mathcal{D}$ and denote it as $\mathcal{Z}=\{z_{1},\cdots,z_{N}\}$.
We only train the weights $W=(w_{1},\cdots,w_{m})\in\mathbb{R}^{m\times d}$ and freeze the values of $a,b$ during the training.
We denote the initial value of the weight and its value at time $t$ as $W^{(0)}$ and $W^{(0)}+W_{t}$ respectively.
Each element of $W^{(0)}$ and $b^{(0)}$ are initialized from $\mathcal{N}(0,1/m)$ while each element of $a^{(0)}_{r}$ are initialized from $\mathcal{N}(0,\varepsilon_{a}^{2})$ for some fixed $\varepsilon_{a}\in(0,1]$.
At each step $t$, we sample a single data point $z=(x,y)$ from $\mathcal{Z}$ and update $W$ using un-normalized version of SAM:

$$ $\displaystyle W_{t+1}$ $\displaystyle=W_{t}-\eta\nabla L(F(x;W^{(0)}+W_{t+1/2}),y)$ $\displaystyle=W_{t}-\eta\nabla L(F(x;W^{(0)}+\rho\nabla L(F(x;W^{(0)}+W_{t}),y )),y).$ (19) $$

### F.2 Formal Theorem

Now, we are ready to present the formal version of Theorem 5.7 below.

###### Theorem F.1 .

(SAM version of Theorem 1 in Allen-Zhu et al. ())
For every $\varepsilon\in\left(0,\frac{1}{pk\mathfrak{C}_{\mathfrak{s}}(\phi,1)}\right)$, there exists $M_{0}=\textsf{{poly}}(\mathfrak{C}_{\epsilon}(\phi,1),1/\varepsilon)$ and $N_{0}=\textsf{{poly}}(\mathfrak{C}_{\mathfrak{s}}(\phi,1),1/\varepsilon)$ such that for every $m\geq M_{0}$ and every $N\geq\widetilde{\Omega}(N_{0})$, by choosing $\varepsilon_{a}=\varepsilon/\widetilde{\Theta}(1)$ for the initialization and $\eta=\widetilde{\Theta}(\frac{1}{\varepsilon km}),\rho=\widetilde{\Theta}(
\frac{1}{\varepsilon^{3}km^{3}}),T=\widetilde{\Theta}\left(\frac{(\mathfrak{C}
_{\mathfrak{s}}(\phi,1))^{2}\cdot k^{3}p^{2}}{\varepsilon^{2}}\right)$, running $T$ iterations of stochastic SAM defined in Section F.1 gives the following generalization bound with high probability over the random initialization.

$$ $\mathbb{E}_{\text{SAM}}\left[\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}_{(x,y)\sim \mathcal{D}}L(F(x;W^{(0)}+W_{t}),y)\right]\leq\textsf{{OPT}}+\varepsilon.$ (20) $$

Here, the notation of $\widetilde{O}(\cdot)$ ignores the factor of $\textsf{polylog}(m)$.

###### Theorem F.1 .

### F.3 Proof of Theorem F.1

We here present the proof of Theorem F.1.

First, note that we can directly use the algorithm-independent part from Allen-Zhu et al. ().
Thus, it is sufficient to show that the similar version of Lemma B.4 in Allen-Zhu et al. () also holds for SAM.

We first define the function $G=(g_{1},\cdots,g_{k}):\mathbb{R}^{d}\rightarrow\mathbb{R}^{k}$ as similar to Allen-Zhu et al. ().

$$ $g_{r}(x;W_{t})\stackrel{{\scriptstyle\rm def}}{{=}}\sum_{i=1}^{m}a^{(0)}_{r,i} (\langle w^{(t)}_{i},x\rangle+b^{(0)}_{i})\mathbbm{1}[\langle w^{(0)}_{i},x \rangle+b^{(0)}_{i}\geq 0].$ (21) $$

Then, the following corollary holds for a stochastic SAM from Lemma B.3 of Allen-Zhu et al. ().
The corollary presents an upper bound on the norm of differences between $\frac{\partial}{\partial W}L(F(\cdot),y)$ and $\frac{\partial}{\partial W}L(G(\cdot),y)$.

###### Corollary F.2 .

(SAM version of Lemma B.3 in Allen-Zhu et al. ())
Let $\tau=\varepsilon_{a}(\eta+\rho)t$.
Then, for every $x$ satisfying $\|x\|_{2}=1$, and for every time step $t\geq 1$, the following are satisfied with high probability over the random initialization.
(a) For every $r\in[k]$,

$$ $\left\lvert f_{r}(x;W^{(0)}+W_{t})-g_{r}(x;W^{(0)}+W_{t})\right\rvert= \widetilde{O}(\varepsilon_{a}k\tau^{2}m^{3/2})$ $$

(b) For every $y\in\mathcal{Y}$,

$$ $\left\|\frac{\partial}{\partial W}L(F(x;W^{(0)}+W_{t}),y)-\frac{\partial}{ \partial W}L(G(x;W^{(0)}+W_{t}),y)\right\|_{2,1}\leq\widetilde{O}(\varepsilon_ {a}k\tau m^{3/2}+\varepsilon_{a}^{2}k^{2}\tau^{2}m^{5/2})$ (22) $$

Next, we present the key lemma integral to our proof.
The part $(c)$ will be directly used in the proof and presents an upper bound on the norm of differences between SAM gradient and SGD gradient for $F$.

###### Lemma F.3 .

For every $x$ satisfying $\|x\|_{2}=1$, and for every time step $t\geq 1$, the following are satisfied with high probability over the random initialization.
(a) For at most $\widetilde{O}(\varepsilon_{a}\rho\sqrt{km})$ fraction of $i\in[m]$: we have

$$ $\mathbbm{1}[\langle w^{(t+1/2)}_{i},x\rangle+b^{(0)}_{i}\geq 0]\neq\mathbbm{1} [\langle w^{(t)}_{i},x\rangle+b^{(0)}_{i}\geq 0].$ $$

(b) For every $r\in[k]$,

$$ $\left\lvert f_{r}(x;W^{(0)}+W_{t+1/2})-f_{r}(x;W^{(0)}+W_{t})\right\rvert= \widetilde{O}(\varepsilon_{a}^{3}k\rho^{2}m^{3/2}+\varepsilon_{a}^{2}\sqrt{k} \rho m)$ $$

(c) For every $y\in\mathcal{Y}$,

$$ $\displaystyle\left\|\frac{\partial}{\partial W}L(F(x;W^{(0)}+W_{t+1/2}),y)- \frac{\partial}{\partial W}L(F(x;W^{(0)}+W_{t}),y)\right\|_{2,1}$ $\displaystyle\leq\widetilde{O}(\varepsilon_{a}^{2}k\rho m^{3/2}+\varepsilon_{a }^{4}k^{2}\rho^{2}m^{5/2}+\varepsilon_{a}^{3}k^{3/2}\rho m^{2})$ (23) $$

###### Proof.

Recall that the following hold from the definition of $F$ (see Lemma B.3 of Allen-Zhu et al. () for the details).

$$ $\left\|\frac{\partial}{\partial w_{i}}f_{r}(x;W^{(0)}+W_{t})\right\|_{2}\leq \varepsilon_{a}B\quad\text{and}\quad\left\|\frac{\partial}{\partial w_{i}}L(F( x;W^{(0)}+W_{t}),y)\right\|_{2}\leq\sqrt{k}\varepsilon_{a}B$ (24) $$

(a)
Let $\tau=\varepsilon_{a}\rho$ and define $\mathcal{H}\stackrel{{\scriptstyle\rm def}}{{=}}\left\{i\in[m]\bigg{\|}\left
\lvert\langle w^{(t)}_{i},x\rangle+b^{(0)}_{i}\right\rvert\geq 2\sqrt{k}B\tau\right\}$.
Then, the lemma is a direct corollary from Lemma B.3 (a) of Allen-Zhu et al. ().
(b)
We divide $i$ into two cases.
First, when $i\notin\mathcal{H}$, we can directly utilize Lemma B.3.(b) of Allen-Zhu et al. () and the total difference from these $i$’s is $\widetilde{O}(\varepsilon_{a}^{3}k\rho^{2}m^{3/2})$.
Next, we consider the differences from $i\in\mathcal{H}$.

$$ $\displaystyle\quad\left\lvert a^{(0)}_{r,i}\left(\left\langle w^{(t+1/2)}_{i}, x\right\rangle+b^{(0)}_{i}\right)\mathbbm{1}\left[\left\langle w^{(t+1/2)}_{i} ,x\right\rangle+b^{(0)}_{i}\geq 0\right]\right.$ $\displaystyle\hskip 20.00003pt\left.-a^{(0)}_{r,i}\left(\left\langle w^{(t)}_{ i},x\right\rangle+b^{(0)}_{i}\right)\mathbbm{1}\left[\left\langle w^{(t)}_{i}, x\right\rangle+b^{(0)}_{i}\geq 0\right]\right\rvert$ $\displaystyle\leq\left\lvert a^{(0)}_{r,i}\left(\left\langle w^{(t+1/2)}_{i}-w ^{(t)}_{i},x\right\rangle\right)\right\rvert$ $\displaystyle=\left\lvert a^{(0)}_{r,i}\left(\left\langle\rho\cdot\frac{ \partial}{\partial w_{i}}L(F(x;W^{(0)}+W_{t}),y),x\right\rangle\right)\right\rvert$ $\displaystyle\leq\rho\left\lvert a^{(0)}_{r,i}\right\rvert\cdot\left\|\frac{ \partial}{\partial w_{i}}L(F(x;W^{(0)}+W_{t}),y)\right\|_{2}\cdot\|x\|_{2}$ $\displaystyle\leq\rho(\varepsilon_{a}B)\cdot(\sqrt{k}\varepsilon_{a}B)$ $\displaystyle=\widetilde{O}(\varepsilon_{a}^{2}\sqrt{k}\rho)$ $$

The first inequality is from the fact that $i\in\mathcal{H}$ and thus $\mathbbm{1}\left[\left\langle w^{(t+1/2)}_{i},x\right\rangle+b^{(0)}_{i}\geq 0
\right]=\mathbbm{1}\left[\left\langle w^{(t)}_{i},x\right\rangle+b^{(0)}_{i}
\geq 0\right]$.
Then, we have utilized the definition of SAM (F.1) and Cauchy-Schwartz inequality.
Since there can be at most $m$ number of $i\in\mathcal{H}$, the total differences from $i\in\mathcal{H}$ amount to $\widetilde{O}(\varepsilon_{a}^{2}\sqrt{k}\rho m)$.
Combining the two cases proves the (b).
(c)
By the chain rule, we have

$$ $\frac{\partial}{\partial w_{i}}L(F(x;W^{(0)}+W_{t}),y)=\nabla L(F(x;W^{(0)}+W_ {t}),y)\frac{\partial}{\partial w_{i}}F(x;W^{(0)}+W_{t}).$ $$

Since $L$ is $1$-smooth, applying the above lemma (b) gives

$$ $\displaystyle\hskip 11.00008pt\left\|\nabla L(F(x;W^{(0)}+W_{t+1/2}),y)-\nabla L (F(x;W^{(0)}+W_{t}),y)\right\|_{2}$ $\displaystyle\leq\left\|F(x;W^{(0)}+W_{t+1/2})-F(x;W^{(0)}+W_{t})\right\|_{2}$ $\displaystyle\leq\widetilde{O}\left(\varepsilon_{a}^{3}k^{3/2}\rho^{2}m^{3/2}+ \varepsilon_{a}^{2}k\rho m\right).$ (25) $$

For $i\in\mathcal{H}$, we have $\mathbbm{1}[\langle w^{(t+1/2)}_{i},x\rangle+b^{(0)}_{i}\geq 0]=\mathbbm{1}[
\langle w^{(t)}_{i},x\rangle+b^{(0)}_{i}\geq 0]$ and thus $\frac{\partial}{\partial w_{i}}F(x;W^{(0)}+W_{t+1/2})=\frac{\partial}{\partial
w
_{i}}F(x;W^{(0)}+W_{t})$.
Then, combining (F.3) with (24) and using the fact that there can be at most $m$ number of $i\in\mathcal{H}$, this amounts to $\widetilde{O}\left(\varepsilon_{a}^{4}k^{2}\rho^{2}m^{5/2}+\varepsilon_{a}^{3}
k^{3/2}\rho m^{2}\right)$.

Next, for $i\notin\mathcal{H}$, we can directly use the result from Lemma B.3.(c) of Allen-Zhu et al. () and this contributes to $\widetilde{O}(\varepsilon_{a}^{2}k\rho m^{3/2})$.
Summing these together, we prove the bound.
∎

Finally, we show that the following lemma holds, which is a SAM version of Lemma B.4 in Allen-Zhu et al. ().
Combined with the algorithm-independent parts presented in Allen-Zhu et al. (), proving the following lemma concludes the proof of Theorem F.1.
We use the notation of $L_{F}(\mathcal{Z};W)$ for $L_{F}(\mathcal{Z};W)\stackrel{{\scriptstyle\rm def}}{{=}}\frac{1}{|\mathcal{Z}
|}\sum_{(x,y)\in\mathcal{Z}}L(F(x;W+W^{(0)}),y)$ and similarly define $L_{G}(\mathcal{Z};W)$.

###### Lemma F.4 .

(SAM version of Lemma B.4 in Allen-Zhu et al. ())
For every $\varepsilon\in\left(0,\frac{1}{pk\mathfrak{C}_{\mathfrak{s}}(\phi,1}\right)$, letting $\varepsilon_{a}=\varepsilon/\widetilde{\Theta}(1)$, $\eta=\widetilde{\Theta}(\frac{1}{\varepsilon km})$, and $\rho=\widetilde{\Theta}(\frac{1}{\varepsilon^{3}km^{3}})$, there exists $M=\textsf{{poly}}(\mathfrak{C}_{\epsilon}(\phi,1),1/\varepsilon)$ and $T=\Theta\left(\frac{k^{3}p^{2}\cdot\mathfrak{C}_{\mathfrak{s}}(\phi,1)^{2}}{
\varepsilon^{2}}\right)$
such that if $m\geq M$, the following holds with high probability over random initialization.

$$ $\frac{1}{T}\sum_{t=0}^{T-1}L_{F}(\mathcal{Z},W_{t})\leq\textsf{{OPT}}+\varepsilon.$ (26) $$

###### Proof.

Let $W^{\text{\textreferencemark}}$ be the weights constructed from the Corollary B.2 in Allen-Zhu et al. ().
By the convexity of $L$ and Cauchy-Schwartz inequality, we have

$$ $\displaystyle L_{G}(\mathcal{Z},W_{t})-L_{G}(\mathcal{Z};W^{\text{ \textreferencemark}})$ $\displaystyle\leq\langle\nabla L_{G}(\mathcal{Z};W_{t}),W_{t}-W^{\text{ \textreferencemark}}\rangle$ $\displaystyle=\langle\nabla L_{G}(\mathcal{Z};W_{t})-\nabla L_{F}(\mathcal{Z}; W_{t}),W_{t}-W^{\text{\textreferencemark}}\rangle$ $\displaystyle\quad+\langle\nabla L_{F}(\mathcal{Z};W_{t})-\nabla L_{F}( \mathcal{Z};W_{t+1/2}),W_{t}-W^{\text{\textreferencemark}}\rangle$ $\displaystyle\quad+\langle\nabla L_{F}(\mathcal{Z};W_{t+1/2}),W_{t}-W^{\text{ \textreferencemark}}\rangle$ $\displaystyle\leq\|\nabla L_{G}(\mathcal{Z};W_{t})-\nabla L_{F}(\mathcal{Z};W_ {t})\|_{2,1}\|W_{t}-W^{\text{\textreferencemark}}\|_{2,\infty}$ $\displaystyle\quad+\|\nabla L_{F}(\mathcal{Z};W_{t})-\nabla L_{F}(\mathcal{Z}; W_{t+1/2})\|_{2,1}\|W_{t}-W^{\text{\textreferencemark}}\|_{2,\infty}$ $\displaystyle\quad+\langle\nabla L_{F}(\mathcal{Z};W_{t+1/2}),W_{t}-W^{\text{ \textreferencemark}}\rangle$ $$

From the SAM update rule (F.1), we have the following equality.

$$ $\displaystyle\|W_{t+1}-W^{\text{\textreferencemark}}\|_{F}^{2}$ $\displaystyle=\|W_{t}-\eta\nabla L_{F}(z^{(t)},W_{t+1/2})-W^{\text{ \textreferencemark}}\|_{F}^{2}$ $\displaystyle=\|W_{t}-W^{\text{\textreferencemark}}\|_{F}^{2}-2\eta\langle \nabla L_{F}(z^{(t)},W_{t+1/2}),W_{t}-W^{\text{\textreferencemark}}\rangle+ \eta^{2}\|\nabla L_{F}(z^{(t)},W_{t+1/2})\|_{F}^{2}.$ $$

Thus, we have

$$ $\displaystyle L_{G}(\mathcal{Z};W_{t})-L_{G}(\mathcal{Z};W^{\text{ \textreferencemark}})$ $\displaystyle\leq\underbrace{\|\nabla L_{G}(\mathcal{Z};W_{t})-\nabla L_{F}( \mathcal{Z}:W_{t})\|_{2,1}\|W_{t}-W^{\text{\textreferencemark}}\|_{2,\infty}}_ {(A)}$ $\displaystyle\quad+\underbrace{\|\nabla L_{F}(\mathcal{Z};W_{t})-\nabla L_{F}( \mathcal{Z};W_{t+1/2})\|_{2,1}\|W_{t}-W^{\text{\textreferencemark}}\|_{2, \infty}}_{(B)}$ $\displaystyle\quad+\frac{\|W_{t}-W^{\text{\textreferencemark}}\|_{F}^{2}- \mathbb{E}_{z^{(t)}}[\|W_{t+1}-W^{\text{\textreferencemark}}\|_{F}^{2}]}{2\eta}$ $\displaystyle\quad+\underbrace{\frac{\eta}{2}\|\nabla L_{F}(W_{t+1/2},z^{(t)}) \|_{F}^{2}}_{(C)}.$ $$

Since $\|W_{t}-W^{\text{\textreferencemark}}\|_{2,\infty}=\widetilde{O}(\sqrt{k}
\varepsilon_{a}(\eta+\rho)t+\frac{kpC_{0}}{\varepsilon_{a}m})$, $(A)$ is bounded as

$$ $(A)=\widetilde{O}\left(\sqrt{k}\varepsilon_{a}(\eta+\rho)T\Delta+\frac{kpC_{0} }{\varepsilon_{a}m}\Delta\right)$ $$

where $\Delta=\widetilde{O}\left(\varepsilon_{a}^{2}k(\eta+\rho)Tm^{3/2}+\varepsilon_
{a}^{4}k^{2}(\eta+\rho)^{2}T^{2}m^{5/2}\right)$.

Next, we can bound $(B)$ from Lemma F.3(c) as follows.

$$ $(B)=\widetilde{O}(\sqrt{k}\varepsilon_{a}(\eta+\rho)T\Delta^{\prime}+\frac{kpC _{0}}{\varepsilon_{a}m}\Delta^{\prime}),$ $$

where $\|\nabla L_{F}(\mathcal{Z};W_{t})-\nabla L_{F}(\mathcal{Z};W_{t+1/2})\|_{2,1}
\leq\Delta^{\prime}=\varepsilon_{a}^{2}k\rho m^{3/2}+\varepsilon_{a}^{4}k^{2}
\rho^{2}m^{5/2}+\varepsilon_{a}^{3}k^{3/2}\rho m^{2}$.

We also have

$$ $(C)=\widetilde{O}(\eta\varepsilon_{a}^{2}km)$ $$

since the norm of $\nabla L_{F}$ is always bounded as $\|\nabla L_{F}(\cdot,z^{(t)})\|_{F}^{2}=\widetilde{O}(\varepsilon_{a}^{2}km)$.

Then, by telescoping, we have

$$ $\displaystyle\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}_{\text{SAM}}[L_{G}(\mathcal {Z};W_{t})]-L_{G}(\mathcal{Z};W^{\text{\textreferencemark}})$ $\displaystyle\leq\widetilde{O}\left(\sqrt{k}\varepsilon_{a}(\eta+\rho)T\Delta+ \frac{kpC_{0}}{\varepsilon_{a}m}\Delta\right)$ $\displaystyle+\widetilde{O}\left(\sqrt{k}\varepsilon_{a}(\eta+\rho)T\Delta^{ \prime}+\frac{kpC_{0}}{\varepsilon_{a}m}\Delta^{\prime}\right)$ $\displaystyle\quad+\underbrace{\frac{\|W_{0}-W^{\text{\textreferencemark}}\|_{ F}^{2}}{2\eta T}}_{(D)}+\widetilde{O}(\eta\varepsilon_{a}^{2}km).$ $$

We can bound $(D)$ in the same way as Allen-Zhu et al. (),

$$ $(D)=\frac{\|W_{0}-W^{\text{\textreferencemark}}\|_{F}^{2}}{2\eta T}=\widetilde {O}\left(\frac{k^{2}p^{2}\mathfrak{C}_{\mathfrak{s}}(\phi,1)^{2}}{\varepsilon_ {a}^{2}m}\cdot\frac{\bm{1}}{\eta T}\right).$ $$

By setting $\eta=\widetilde{\Theta}(\frac{\varepsilon}{km\varepsilon_{a}^{2}}),\rho=
\widetilde{\Theta}(\frac{\varepsilon}{km^{3}\varepsilon_{a}^{4}}),T=\widetilde
{\Theta}(k^{3}p^{2}\mathfrak{C}_{\mathfrak{s}}(\phi,1)^{2}/\varepsilon^{2})$, we have
$\Delta=\widetilde{O}(\frac{k^{6}p^{4}\mathfrak{C}_{\mathfrak{s}}(\phi,1)^{4}}{
m^{3/2}\varepsilon^{4}})$ and
$\Delta^{\prime}=\widetilde{O}(\frac{1}{m^{3/2}\varepsilon}+\frac{\sqrt{k}}{m})$.
Hence, with large enough $m$, we obtain the following inequality and prove Lemma F.4, combined with the remaining parts from Allen-Zhu et al. ().

$$ $\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}_{\text{SAM}}[L_{G}(\mathcal{Z};W_{t})]-L _{G}(\mathcal{Z};W^{\text{\textreferencemark}})\leq O(\varepsilon).$ (27) $$

∎

###### Corollary F.2 .

###### Lemma F.3 .

###### Proof.

###### Lemma F.4 .

###### Proof.

## Appendix G Proof of Theorem 5.2

Here, we provide the detailed proof of Theorem 5.2.

We first define a linearized stochastic SAM, which is derived from applying first-order Taylor approximation on a stochastic SAM update given as follows:

###### Definition G.1 .

(Linearized stochastic SAM)
We define a linearized stochastic SAM as

$$ $x_{t+1}=x_{t}-\eta H_{\xi_{t}}(x_{t+1/2}-x^{\star}),$ (28) $$

where $x_{t+1/2}=x_{t}+\rho H_{\xi_{t}}(x_{t}-x^{\star})$ is the linearized ascent step and $H_{\xi_{t}}$ is the Hessian estimation at step $t$.

This actually corresponds to using SAM for the quadratic approximation of $f$ near $x^{\star}$, and we use this fact in the experiment setup.
We assume without loss of generality that the fixed point $x^{\star}$ satisfies $x^{\star}=0$.

Then, we are ready to present the proof of Theorem 5.2. Our goal is to derive a bound of the form $\mathbb{E}\|x_{t}\|^{2}\leq C\|x_{0}\|^{2}$.
We first apply (28) to $\mathbb{E}{\left[\|x_{t+1}\|^{2}\,\lvert\,x_{t}\right]}$ and continue expanding the terms as follows:

$$ $\displaystyle\mathbb{E}{\left[\|x_{t+1}^{2}\|\,\lvert\,x_{t}\right]}$ $\displaystyle=\mathbb{E}\|x_{t}-\eta H_{\xi_{t}}(x_{t}+\rho H_{\xi_{t}}x_{t}) \|^{2}$ $\displaystyle=x_{t}^{\top}\mathbb{E}{\left[\left(I-\eta H_{\xi_{t}}-\eta\rho H _{\xi_{t}}^{2}\right)^{2}\,\Big{\lvert}\,x_{t}\right]}x_{t}$ $\displaystyle=x_{t}^{\top}\mathbb{E}{\left[I-2\eta(H_{\xi_{t}}+\rho H_{\xi_{t} }^{2})+\eta^{2}\left(H_{\xi_{t}}+\rho H_{\xi_{t}}^{2}\right)^{2}\,\Big{\lvert} \,x_{t}\right]}x_{t}$ $\displaystyle=x_{t}^{\top}\mathbb{E}{\left[I-2\eta(H_{\xi_{t}}+\rho H_{\xi_{t} }^{2})+\eta^{2}\left(H_{\xi_{t}}^{2}+2\rho H_{\xi_{t}}^{3}+\rho^{2}H_{\xi_{t}} ^{4}\right)\,\Big{\lvert}\,x_{t}\right]}x_{t}$ $\displaystyle=x_{t}^{\top}\mathbb{E}{\left[I-2\eta H_{\xi_{t}}+\eta(\eta-2\rho )H_{\xi_{t}}^{2}+2\eta^{2}\rho H_{\xi_{t}}^{3}+\eta^{2}\rho^{2}H_{\xi_{t}}^{4} \,\Big{\lvert}\,x_{t}\right]}x_{t}$ $\displaystyle=x_{t}^{\top}\Big{(}I-2\eta H+\eta(\eta-2\rho)\mathbb{E}H_{\xi_{t }}^{2}+2\eta^{2}\rho\mathbb{E}H_{\xi_{t}}^{3}+\eta^{2}\rho^{2}\mathbb{E}H_{\xi _{t}}^{4}\Big{)}x_{t}$ $\displaystyle=x_{t}^{\top}\Big{(}I-2\eta H+\eta(\eta-2\rho)H^{2}+2\eta^{2}\rho H ^{3}+\eta^{2}\rho^{2}H^{4}$ $\displaystyle\hskip 30.00005pt+\eta(\eta-2\rho)(\mathbb{E}H_{\xi_{t}}^{2}-H^{2 })+2\eta^{2}\rho(\mathbb{E}H_{\xi_{t}}^{3}-H^{3})+\eta^{2}\rho^{2}(\mathbb{E}H _{\xi_{t}}^{4}-H^{4})\Big{)}x_{t}$ $\displaystyle=x_{t}^{\top}\Big{(}\left(I-\eta H-\eta\rho H^{2}\right)^{2}$ $\displaystyle\hskip 40.00006pt+\eta(\eta-2\rho)(\mathbb{E}H_{\xi_{t}}^{2}-H^{2 })+2\eta^{2}\rho(\mathbb{E}H_{\xi_{t}}^{3}-H^{3})+\eta^{2}\rho^{2}(\mathbb{E}H _{\xi_{t}}^{4}-H^{4})\Big{)}x_{t}$ $$

Since $x^{\top}Ax\leq\lambda_{\text{max}}(A)\|x\|^{2}$ always holds for any $x$ and any matrix $A$ with the maximum eigenvalue $\lambda_{\text{max}}(A)$, applying this inequality and taking the total expectation gives the following;

$$ $\displaystyle\mathbb{E}{\left[\|x_{t+1}\|^{2}\right]}$ $\displaystyle\leq\lambda_{\text{max}}\bigg{(}\left(I-\eta H-\eta\rho H^{2} \right)^{2}+\eta(\eta-2\rho)(\mathbb{E}H_{\xi}^{2}-H^{2})$ $\displaystyle\hskip 50.00008pt+2\eta^{2}\rho(\mathbb{E}H_{\xi}^{3}-H^{3})+\eta ^{2}\rho^{2}(\mathbb{E}H_{\xi}^{4}-H^{4})\bigg{)}\mathbb{E}{\left[\|x_{t}\|^{2 }\right]}.$ $$

Recursively applying this bound gives

$$ $\displaystyle\mathbb{E}\|x_{t}\|^{2}$ $\displaystyle\leq\lambda_{\text{max}}\bigg{(}\left(I-\eta H-\eta\rho H^{2} \right)^{2}+\eta(\eta-2\rho)(\mathbb{E}H_{\xi}^{2}-H^{2})$ $\displaystyle\hskip 50.00008pt+2\eta^{2}\rho(\mathbb{E}H_{\xi}^{3}-H^{3})+\eta ^{2}\rho^{2}(\mathbb{E}H_{\xi}^{4}-H^{4})\bigg{)}^{t}\|x_{0}\|^{2}.$ $$

Here, we can see that $x^{\star}$ is linearly stable if

$$ $\displaystyle\lambda_{\text{max}}\bigg{(}(I-\eta H-\eta\rho H^{2})^{2}$ $\displaystyle\hskip 40.00006pt+\eta(\eta-2\rho)(\mathbb{E}H_{\xi}^{2}-H^{2})+2 \eta^{2}\rho(\mathbb{E}H_{\xi}^{3}-H^{3})+\eta^{2}\rho^{2}(\mathbb{E}H_{\xi}^{ 4}-H^{4})\bigg{)}\leq 1.$ $$

###### Definition G.1 .

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 6
presentation: 4
questions: This is an interesting paper. I have a question: how does the convergence of higher-order moments of Hessian in your result compare with the other approaches in the literature, e.g., [1]? Can you provide a literature review on the previous works considering higher-order moments of Hessian to define flatness?

[1] Tahmasebi, Behrooz, et al. "A Universal Class of Sharpness-Aware Minimization Algorithms." Forty-first International Conference on Machine Learning.
soundness: 3
strengths: - interesting and well-motivated problem
 - very well written
summary: The authors perform experiments to 
measure the effect of overparameterization in SAM for a diverse set of tasks (Section 3). The goal is to observe how overparameterization
affects SAM under various conditions, e.g.,  label noise, sparsity, and regularization. 
Furthermore, they prove that stable minima of SAM are flatter and have more uniform Hessian moments (if compared with SGD), and stochastic SAM can also converge at a linear. The overall contribution is that they empirically and theoretically proved that overparameterization critically affects SAM.
weaknesses: - discussion on higher moments of Hessian is missing

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Could the authors provide adequate theoretical support for the main findings? e.g., an analysis on diagonal linear networks.
soundness: 2
strengths: The experiments in this paper are insightful, showing that overparameterization can increase the performance gap between SAM and SGD, while also highlighting the role of factors such as label noise, sparsity, weight decay, and early stopping in this phenomenon.
summary: This paper studies the critical influence of overparameterization on SAM.
weaknesses: - It is widely recognized that overparameterization improves the generalization performance of SGD (e.g., ResNet-152 outperforms ResNet-18 on Cifar10 using SGD). And it is unsurprising that overparameterization can also improve the generalization performance of SAM. However, the non-trivial aspect is that overparameterization increases *the gap between SAM and SGD*, as shown in Figure 1. To avoid trivialization, reconsidering the title and abstract might be beneficial.

- My primary concern is that the theoretical analysis can not support the main claim/findings (e.g., Figure 1). 
  - The linear stability analysis in Section 6.1 can not demonstrate how overparameterization affects SAM (i.e., that “greater overparameterization leads SAM to find flatter minima”).
  - In the convergence analysis in Section 6.2: (i) it lacks relevance to generalization, the core focus of the paper; (ii) it does not clarify how the degree of overparameterization influences the convergence speed within the interpolation regime; (iii) in the interpolation regime, the exponential convergence result also holds for SGD (Bassily et al., 2018). The proof for SAM appears to be a minor extension from SGD, which treats SAM as SGD + small permutation.
  - The theoretical analysis in Section 6.1/6.2 focus on a variant of SAM (SAM without normalization) for simplification, rather than the original SAM. While not a major issue, the formulation should be clearly presented in Section 6 to prevent potential reader confusion.

---

confidence: 2
contribution: 1
initial_rating: 6
presentation: 3
questions: N/A
soundness: 2
strengths: 1. The authors picked up an interesting perspective to investigate the mechanism of SAM. This research topic is significant.
2. The experiments are extensive, across various datasets and architectures.
3. The writing structure is clear, and the paper is easy to follow.
summary: This paper claimed that SAM benefits from the over-parameterization. They find that SAM could find different (e.g. simpler) solutions than SGD in the over-parameterization regime. In addition, they find that SAM is robust to label noise and sparse parameterization. They also include some theoretical analyses, thought these analyses are irrelevant to their main findings.
weaknesses: Firstly, for the empirical findings,

1. It is unclear that whether over-parameterization consistently helps SAM find more generalizable solutions as the model parameters scale up. As shown in Fig.1 and Fig. 7 (I think the Fig. 7 is clearer), for MLP on MNIST, ResNet-18 on CIFAR-10, LSTM on SST2 and CNN on Atari, when model get more over-parameterized, the generalization improvement does not always increase; sometimes, a noticeable decline is observed. Moreover, as noted by the authors, when there is no sufficient regularization, SAM might not benefit from the over-parameterization (as shown in Fig. 5). In addition, I notice that the authors conducted experiments for ViT on CIFAR-10 in Fig.5 but not in Fig. 1. I wish the authors could include the results for ViT on CIFAR-10 in Fig.1.
2. Increasing depth of models yields different conclusions than increasing width. In Fig. 1, authors increase the number of neurons in each layer, leading to a seemingly consistent conclusion that over-parameterization helps SAM find generalizable solutions. However, in Fig. 17, when increasing the number of layers, SAM might not benefit from the over-parameterization (ResNet on CIFAR-10). Clearly, the increase of both depth and width demonstrates the over-parameterization, however, the conclusion differs. 
3. It is not surprising that the marginal improvement of SAM over SGD increases with higher label noise. In Fig. 12, in the over-parameterization regime, the performance of SGD drops significantly with higher label noise rate; however, SAM remains good generalization ability. Indeed, it is well known that SAM is robust to label noise, while SGD not. Thus, it is not surprising for us to observe such phenomenon. Also, it seems that the authors over-claimed that over-parameterization secures the robustness of SAM to label noise. Indeed, the observed increase of margin improvement of SAM over SGD with larger model parameters is primarily due to the pronounced decline in SGD's performance as the model parameters scale up under high label noise conditions.
4. The conclusion on the effect of sparsity contradicts with the main argument that SAM benefits from the over-parameterization. In Fig. 5, it is observed that the generalization improvement tends to increase as the model becomes sparser. However, increased sparsity indicates that the model is becoming more under-parameterized. The authors are suggested to explain such contradiction.
5. In Fig. 2 and 3, the authors compare GD and SAM, however, a comparison between SGD and SAM is more preferred.
6. In section 4.2, “implicit bias of SAM increase with over-parameterization” is over-stated. In this section, the authors demonstrated that we could use larger perturbation radius when scaling up the model parameters, however, this doesn’t indicate that a larger scale of model parameters induces stronger implicit bias towards flatter minima. 

In short, SAM doesn’t always benefit from the over-parameterization in practice. A more thorough justification is needed to clarify whether over-parameterization is advantageous for SAM and under what specific conditions it may provide benefits. Indeed, it is quite expectable that SAM would benefit from over-parameterization sometimes. Because in the under-parameterization regime, the model capacity is low (say only one solution in the solution space), thus both SGD and SAM achieve similar solutions. Once scaling up the parameters, the solution space is getting larger, thereby SAM could differ from SGD and find more generalizable solutions. However, SAM might not always from the over-parameterization. 

Second, for the theoretical analysis,

1. As noted by the authors, the theoretical analysis cannot support the findings in Section 3 and 4, and thus the significance of the theoretical analysis is limited. The section 6 is more like a collection of possible theoretical analyses that could be done for SAM. 
2. Indeed, there is no clear theoretical support for the main discovery of this paper, as noted by the authors in the limitations part in Section 7. This is another critical issue of this paper.

Overall, the major issue of this paper is the novelty. The empirical findings are largely anticipated, and the theoretical analysis closely follows prior works without introducing any new proof techniques.  Also, some statements are clearly over-claimed. The contribution of this paper is marginal, and thus I lean to rejection.


----------------
As I cannot respond in official comments further, I put my responses here:

Based on your clarification on the ViT experiment, I decide to increase my score from 5 to 6. However, I strongly recommend the authors to revise the paper from the following aspects:

- **Overstated claims**: Some statements, like "over-parameterization consistently improves the generalization benefit of SAM" in the introduction part, are over-stated, which might be misleading for the readers. Some failure experiments such as ViT, increasing depth of networks, are encouraged to put earlier in the main text, along with throughout discussions. *I firmly believe that adhering to strictness is not a bad thing for a scientific paper*.
- **Experiments on Absolute validation metrics**: Experiments using the absolute validation metrics are encouraged to put into main text along with the experiments using "Acc improvements".
- **Current theory cannot directly support your main finding**: Still, currently, your theory cannot directly support your finding. I am not diminishing your contribution, but clearly those complementary theories in current paper cannot persuade me to highly evaluate your paper. I wish the authors could consider my suggestions and improve the paper with more related theoretical analyses. 

BTW, I decide to keep my confidence score low.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: (1) Please clarify the large curve surrounding the "real" curve in Fig . 1 (and possibly in other Figs). Specify if the shaded regions represent confidence intervals, standard deviations or soe other measures of uncertainty. 
(2) How can SAM improvements be compared to SGD improvements via overparametrisation? Hoa SAM and SGD performance changes with increasing model size (may be you can add provide (in Annex) plots showing relative improvement over a baseline for both optimizers across different nodel sizes)
(3) Min Max optimization has some limitations (High complexity; Oscillatory behavior, sensitivity to initialization, etc.), how these problems are addressed here ? Can equation (3) (or some of its properties) help mitigate these issues ? Did you observed these specific issues in your experiments and if so, how have you addressed them ?
(4) Can we use Elliptic Loss Regularisation for SAM instead of Min Max? Please comment on whether Elliptic Loss Regularization could be a viable alternative to the min-max formulation in SAM, and what potential advantages or disadvantages it might have in this context.
(5) Is there a sort of "optimal" number of parameters for which the computation price is acceptable ?Did you  observed any diminishing returns in performance improvement as model size increased. Please provide guidance on balancing computational cost with performance gains when using SAM in practice
soundness: 3
strengths: Nice empirical and theoretical work on overparametrisation on sharpness-aware minimization. It is interesting and useful for ICLR communauty
summary: This work attempts to disclose the critical influence of overparamatrisation on  Sharpness-Aware Minimization  (SAM) from experimental and theoretical perspectives. It starts with an extensive evaluation to display a  highly consistent view showing the generalization benefit of SAM increases with overparametrisation. Without such overparametrisation  SAM may not work. This leads authors to consider the benefit in terms of more solutions and implicit bias. Finally, they developed theoretical advantages and advantages of overparametrization for SAM on linear stability, convergence and generalization.
weaknesses: See questions below.

Meta-review: This paper examines the impact of overparameterization on sharpness-aware minimization, focusing on factors such as label noise, sparsity, and regularization. The reviewers found the topic to be well-motivated and the experiments extensive. However, major concerns were raised regarding the empirical and theoretical grounding of key claims, which were at times perceived as overstated or predictable. Despite the significant discussions the paper has generated, the initial evaluations remained lukewarm. The authors are encouraged to incorporate the important feedback given by the knowledgeable reviewers.

Decision: Reject