SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \gamma_ {j, r} ^ {(t, b + 1)} = \gamma_ {j, r} ^ {(t, b)} - \frac {\eta}{B m} \cdot \left[ \sum_ {i \in \mathcal {I} _ {t, b} \cap S _ {+}} \ell_ {i} ^ {\prime (t, b)} \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \widehat {y} _ {i} \cdot \boldsymbol {\mu} \rangle) \right. \\ \left. - \sum_ {i \in \mathcal {I} _ {t, b} \cap S _ {-}} \ell_ {i} ^ {\prime (t, b)} \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j, r} ^ {(t, b)}, \widehat {y} _ {i} \cdot \boldsymbol {\mu} \right\rangle\right) \right] \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2}. \tag {8} \\ \end{array}
$$
Here,  $\mathcal{I}_{t,b}$  represents the indices of samples in batch  $b$  of epoch  $t$ ,  $S_{+}$  denotes the set of clean samples where  $y_{i} = \widehat{y}_{i}$ , and  $S_{-}$  represents the set of noisy samples where  $y_{i} = -\widehat{y}_{i}$ . The updates of  $\gamma_{j,r}^{(t,b)}$  comprise an increment arising from sample learning, counterbalanced by a decrement due to noisy sample learning. Both empirical and theoretical analyses have demonstrated that overparametrization allows the model to fit even random labels. This occurs when the negative term  $\sum_{i\in \mathcal{I}_{t,b}\cap S_{-}}\ell_i^{\prime (t,b)}\sigma '(\langle \mathbf{w}_{j,r}^{(t,b)},y_i\cdot \boldsymbol {\mu}\rangle)$  primarily drives model learning. Such unfavorable scenarios can be attributed to two possible factors. Firstly, the gradient of the loss  $\ell_{i}^{\prime (t,b)}$  might be significantly larger for noisy samples compared to clean samples. Secondly, during certain epochs, the majority of samples may be noisy, meaning that  $\mathcal{I}_{t,b}\cap S_{-}$  significantly outnumbers  $\mathcal{I}_{t,b}\cap S_{+}$ .
To deal with the first factor, we have to control the ratio of the loss gradient with regard to different samples, as depicted in (9). Given that noisy samples may overwhelm a single batch, we impose an additional requirement: the ratio of the loss gradient must be controllable across different batches within a single epoch, i.e.,
$$
\ell_ {i} ^ {\prime (t, b _ {1})} / \ell_ {k} ^ {\prime (t, b _ {2})} \leq C _ {2}. \tag {9}
$$
As  $\ell'(z_1) / \ell'(z_2) \approx \exp(z_2 - z_1)$ , we can upper bound  $\ell_i'^{(t,b_1)} / \ell_k'^{(t,b_2)}$  by  $y_i \cdot f(\mathbf{W}^{(t,b_1)}, \mathbf{x}_i) - y_k \cdot f(\mathbf{W}^{(t,b_2)}, \mathbf{x}_k)$ , which can be further upper bounded by  $\sum_r \overline{\rho}_{y_i, r,i}^{(t,b_1)} - \sum_r \overline{\rho}_{y_i, r,k}^{(t,b_2)}$  with a small error. Therefore, the proof of (9) can be reduced to proving a uniform bound on the difference among  $\overline{\rho}_{y_i, r,i}^{(t,b)}$ , i.e.,  $\sum_{r=1}^m \overline{\rho}_{y_i, r,i}^{(t,b_1)} - \sum_{r=1}^m \overline{\rho}_{y_k, r,k}^{(t,b_2)} \leq \kappa$ ,  $\forall i, k$ .
However, achieving this uniform upper bound turns out to be challenging, since the updates of  $\overline{\rho}_{j,r,i}^{(t,b)}$ 's are not evenly distributed across different batches within an epoch. Each mini-batch update utilizes only a portion of the samples, meaning that some  $\overline{\rho}_{y_i,r,i}^{(t,b)}$  can increase or decrease much more than
the others. Therefore, the uniformly bounded difference can only be achieved after the entire epoch is processed. Consequently, we have to first bound the difference among  $\overline{\rho}_{y_i,r,i}^{(t,b)}$  's after each entire epoch, and then control the maximal difference within one epoch. The full batch (epoch) update rule is established as follows:
$$
\begin{array}{l} \sum_ {r = 1} ^ {m} \left[ \bar {\rho} _ {y _ {i}, r, i} ^ {(t + 1, 0)} - \bar {\rho} _ {y _ {k}, r, k} ^ {(t + 1, 0)} \right] = \sum_ {r = 1} ^ {m} \left[ \bar {\rho} _ {y _ {i}, r, i} ^ {(t, 0)} - \bar {\rho} _ {y _ {k}, r, k} ^ {(t, 0)} \right] - \frac {\eta (P - 1) ^ {2}}{B m} \cdot \left(\left| \widetilde {S} _ {i} ^ {(t , b _ {i} ^ {(t)})} \right| \ell_ {i} ^ {\prime (t, b _ {i} ^ {(t)})} \right. \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2} \\ \left. - \left| \widetilde {S} _ {k} ^ {(t, b _ {k} ^ {(t)})} \right| \ell_ {k} ^ {\prime (t, b _ {k} ^ {(t)})} \cdot \| \boldsymbol {\xi} _ {k} \| _ {2} ^ {2}\right). \tag {10} \\ \end{array}
$$
Here,  $b_{i}^{(t)}$  denotes the batch to which sample  $i$  belongs in epoch  $t$ , and  $\widetilde{S}_{i}^{(t,b_{i}^{(t)})}$  represents the parameters that learn  $\xi_{i}$  at epoch  $t$  defined as
$$
\widetilde {S} _ {i} ^ {(t, b)} := \left\{r: \left\langle \mathbf {w} _ {y _ {i}, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \right\rangle > 0 \right\}. \tag {11}
$$
Therefore, the update of  $\sum_{r=1}^{m}\left[\overline{\rho}_{y_i,r,i}^{(t,0)} - \overline{\rho}_{y_k,r,k}^{(t,0)}\right]$  is indeed characterized by the activation pattern of parameters, which serves as the key technique for analyzing the full batch update of  $\sum_{r=1}^{m}\left[\overline{\rho}_{y_i,r,i}^{(t,0)} - \overline{\rho}_{y_k,r,k}^{(t,0)}\right]$ . However, analyzing the pattern of  $S_i^{(t,b)}$  directly is challenging since  $\langle \mathbf{w}_{y_i,r}^{(t,b)}, \pmb{\xi}_i \rangle$  fluctuates in mini-batches without sample  $i$ . Therefore, we introduce the set series  $S_i^{(t,b)}$  as the activation pattern with certain threshold as follows:
$$
S _ {i} ^ {(t, b)} := \left\{r: \left\langle \mathbf {w} _ {y _ {i}, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \right\rangle > \sigma_ {0} \sigma_ {p} \sqrt {d} / \sqrt {2} \right\}. \tag {12}
$$
The following lemma suggests that the set of activated parameters  $S_{i}^{(t,0)}$  is a non-decreasing sequence with regards to  $t$ , and the set of plain activated parameters  $\widetilde{S}_{i}^{(t,b)}$  always include  $S_{i}^{(t,0)}$ . Consequently,  $S_{i}^{(0,0)}$  is always included in  $\widetilde{S}_{i}^{(t,b)}$ , guaranteeing that  $\xi_{i}$  can always be learned by some parameter. And this further makes sure the difference among  $\overline{\rho}_{y_i,r,i}^{(t,b)}$  is bounded, as well as  $\ell_{i}^{\prime (t,b_{1})} / \ell_{k}^{\prime (t,b_{2})}\leq C_{2}$ . In the proof for SGD, we consider the learning period  $0\leq t\leq T^{*}$ , where  $T^{*} = \eta^{-1}\mathrm{poly}(\epsilon^{-1},d,n,m)$  is the maximum number of admissible iterations.
Lemma 3.4 (Informal Statement of Lemma C.8). For all  $t \in [0, T^*]$  and  $b \in [\overline{H}]$ , we have
$$
S _ {i} ^ {(t - 1, 0)} \subseteq S _ {i} ^ {(t, 0)} \subseteq \widetilde {S} _ {i} ^ {(t, b)}. \tag {13}
$$
As we have mentioned above, if noisy samples outnumber clean samples,  $\gamma_{j,r}^{(t,b)}$  may also decrease. To deal with such a scenario, we establish a two-stage analysis of  $\gamma_{j,r}^{(t,b)}$  progress. In the first stage, when  $-\ell_i^{\prime}$  is lower bound by a positive constant, we prove that there are enough batches containing sufficient clear samples. This is characterized by the following high-probability event.
Lemma 3.5 (Infomal Statement of Lemma B.6). With high probability, for all  $T \in [\widetilde{O}(1), T^{*}]$ , there exist at least  $c_{1} \cdot T$  epochs among  $[0, T]$ , such that at least  $c_{2} \cdot H$  batches in each of these epochs satisfying the following condition:
$$
\left| S _ {+} \cap S _ {y} \cap \mathcal {I} _ {t, b} \right| \in [ 0. 2 5 B, 0. 7 5 B ]. \tag {14}
$$
After the first stage of  $T = \Theta (\eta^{-1}m(P - 1)^{-2}\sigma_p^{-2}d^{-1})$  epochs, we have  $\gamma_{j,r}^{(T,0)} = \Omega \left(n\frac{\|\pmb{\mu}\|_2^2}{(P - 1)^2\sigma_p^2d}\right)$ . The scale of  $\gamma_{j,r}^{(T,0)}$  guarantees that  $\langle \mathbf{w}_{j,r}^{(t,b)},\pmb {\mu}\rangle$  remains resistant to intra-epoch fluctuations. Consequently, this implies the sign of  $\langle \mathbf{w}_{j,r}^{(t,b)},\pmb {\mu}\rangle$  will persist unchanged throughout the entire epoch. Without loss of generality, we suppose that  $\langle \mathbf{w}_{j,r}^{(t,b)},\pmb {\mu}\rangle >0$ , then the update of  $\gamma_{j,r}^{(t,b)}$  can be written as follows:
$$
\gamma_ {j, r} ^ {(t + 1, 0)} = \gamma_ {j, r} ^ {(t, 0)} + \frac {\eta}{B m} \cdot \left[ \min  _ {i \in \mathcal {I} _ {t, b, b}} \left| \ell_ {i} ^ {\prime (t, b)} \right| | S _ {+} \cap S _ {1} | - \max  _ {i \in \mathcal {I} _ {t, b, b}} \left| \ell_ {i} ^ {\prime (t, b)} \right| | S _ {-} \cap S _ {- 1} | \right] \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2}. \tag {15}
$$
As we have proved the balancing of logits  $\ell_i^{\prime (t,b)}$  across batches, the progress analysis of  $\gamma_{j,r}^{(t + 1,0)}$  is established to characterize the signal learning of SGD.
4 Result for SAM
In this section, we present the positive results for SAM in the following theorem.
Theorem 4.1. For any  $\epsilon >0$ , under Condition 3.1 with  $\sigma_0 = \widetilde{\Theta} (P^{-1}\sigma_p^{-1}d^{-1 / 2})$ , choose  $\tau = \Theta \Big(\frac{m\sqrt{B}}{P\sigma_p\sqrt{d}}\Big)$ . With probability at least  $1 - \delta$ , neural networks first trained with SAM with  $O\Big(\eta^{-1}\epsilon^{-1}n^{-1}mB\| \pmb {\mu}\| _2^{-2}\Big)$  iterations, then trained with SGD with  $\widetilde{O}\Big(\eta^{-1}\epsilon^{-1}mnd^{-1}P^{-2}\sigma_p^{-2}\Big)$  iterations can find  $\mathbf{W}^{(t)}$  such that,
1. The training loss satisfies  $L_{S}(\mathbf{W}^{(t)})\leq \epsilon$
2. The test error  $L_{\mathcal{D}}^{0 - 1}(\mathbf{W}^{(t)}) \leq p + \epsilon$ .
In contrast to Theorem 3.2, Theorem 4.1 demonstrates that CNNs trained by SAM exhibit benign overfitting under much milder conditions. This condition is almost dimension-free, as opposed to the threshold of  $\| \pmb{\mu}\|_{2}^{4}\geq \widetilde{\Omega} ((d / n)P^{4}\sigma_{p}^{4})$  for CNNs trained by SGD. The discrepancy in the thresholds can be observed in Figure 1. This difference is because SAM introduces a perturbation during the model parameter update process, which effectively prevents the early-stage memorization of noise by deactivating the corresponding neurons.
4.1 Noise Memorization Prevention
In this subsection, we will show how SAM can prevent noise memorization by changing the activation pattern of the neurons. For SAM, we have the following update rule of decomposition coefficients  $\gamma_{j,r}^{(t,b)},\overline{\rho}_{j,r,i}^{(t,b)},\underline{\rho}_{j,r,i}^{(t,b)}$
Lemma 4.2. The coefficients  $\gamma_{j,r}^{(t,b)},\overline{\rho}_{j,r,i}^{(t,b)},\underline{\rho}_{j,r,i}^{(t,b)}$  defined in Lemma 3.3 satisfy the following iterative equations for all  $r\in [m]$ $j\in \{\pm 1\}$  and  $i\in [n]$  ..
$$
\gamma_ {j, r} ^ {(0, 0)}, \bar {\rho} _ {j, r, i} ^ {(0, 0)}, \underline {{\rho}} _ {j, r, i} ^ {(0, 0)} = 0,
$$
$$
\begin{array}{l} \gamma_ {j, r} ^ {(t, b + 1)} = \gamma_ {j, r} ^ {(t, b)} - \frac {\eta}{B m} \cdot \left[ \sum_ {i \in \mathcal {I} _ {t, b} \cap S _ {+}} \ell_ {i} ^ {\prime (t, b)} \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j, r} ^ {(t, b)} + \widehat {\boldsymbol {\epsilon}} _ {j, r} ^ {(t, b)}, \widehat {y} _ {i} \cdot \boldsymbol {\mu} \right\rangle\right) \right. \\ \left. - \sum_ {i \in \mathcal {I} _ {t, b} \cap S _ {-}} \ell_ {i} ^ {\prime (t, b)} \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {j, r} ^ {(t, b)} + \widehat {\boldsymbol {\epsilon}} _ {j, r} ^ {(t, b)}, \widehat {y} _ {i} \cdot \boldsymbol {\mu} \right\rangle\right) \right] \cdot \| \boldsymbol {\mu} \| _ {2} ^ {2}, \\ \end{array}
$$
$$
\overline {{\rho}} _ {j, r, i} ^ {(t, b + 1)} = \overline {{\rho}} _ {j, r, i} ^ {(t, b)} - \frac {\eta (P - 1) ^ {2}}{B m} \cdot \ell_ {i} ^ {\prime (t, b)} \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)} + \widehat {\boldsymbol {\epsilon}} _ {j, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \rangle) \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2} \cdot \mathbb {1} (y _ {i} = j)   \mathbb {1} (i \in \mathcal {I} _ {t, b}),
$$
$$
\underline {{\rho}} _ {j, r, i} ^ {(t, b + 1)} = \underline {{\rho}} _ {j, r, i} ^ {(t, b)} + \frac {\eta (P - 1) ^ {2}}{B m} \cdot \ell_ {i} ^ {\prime (t, b)} \cdot \sigma^ {\prime} (\langle \mathbf {w} _ {j, r} ^ {(t, b)} + \widehat {\boldsymbol {\epsilon}} _ {j, r} ^ {(t, b)}, \boldsymbol {\xi} _ {i} \rangle) \cdot \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2} \cdot \mathbb {1} (y _ {i} = - j)   \mathbb {1} (i \in \mathcal {I} _ {t, b}),
$$
where  $\mathcal{I}_{t,b}$  denotes the sample index set of the  $b$ -th batch in the  $t$ -th epoch.
The primary distinction between SGD and SAM is how neuron activation is determined. In SAM, the activation is based on the perturbed weight  $\mathbf{w}_{j,r}^{(t,b)} + \widehat{\pmb{\epsilon}}_{j,r}^{(t,b)}$ , whereas in SGD, it is determined by the unperturbed weight  $\mathbf{w}_{j,r}^{(t,b)}$ . This perturbation to the weight update process at each iteration gives SAM an intriguing denoising property. Specifically, if a neuron is activated by noise in the SGD update, it will subsequently become deactivated after the perturbation, as stated in the following lemma.
Lemma 4.3 (Informal Statement of Lemma D.5). Suppose the Condition 3.1 holds with parameter choices in Theorem 4.1, if  $\langle \mathbf{w}_{j,r}^{(t,b)},\pmb{\xi}_k\rangle \geq 0$ ,  $k\in \mathcal{I}_{t,b}$  and  $j = y_{k}$ , then  $\langle \mathbf{w}_{j,r}^{(t,b)} + \widehat{\pmb{\epsilon}}_{j,r}^{(t,b)},\pmb{\xi}_k\rangle < 0$
By leveraging this intriguing property, we can derive a constant upper bound for the noise coefficients  $\overline{\rho}_{j,r,i}^{(t,b)}$  by considering the following cases:
1. If  $\xi_{i}$  is not in the current batch, then  $\overline{\rho}_{j,r,i}^{(t,b)}$  will not be updated in the current iteration.
(a) Original SGD Test Error Heatmap
(b) SAM Test Error Heatmap
Figure 2: (a) is a heatmap illustrating test error on synthetic data for various dimensions  $d$  and signal strengths  $\mu$  when trained using Vanilla Gradient Descent. High test errors are represented in blue, while low test errors are shown in yellow. (b) displays a heatmap of test errors on the synthetic data under the same conditions as in (a), but trained using SAM instead with  $\tau = 0.03$ . The y-axis represents a normal scale with a range of  $1000 \sim 21000$ .
2. If  $\xi_{i}$  is in the current batch, we discuss two cases:
(a) If  $\langle \mathbf{w}_{j,r}^{(t,b)},\pmb{\xi}_i\rangle \geq 0$ , then by Lemma 4.3, one can know that  $\sigma^{\prime}(\langle \mathbf{w}_{j,r}^{(t,b)} + \widehat{\epsilon}_{j,r}^{(t,b)},\pmb{\xi}_i\rangle) = 0$  and thus  $\overline{\rho}_{j,r,i}^{(t,b)}$  will not be updated in the current iteration.
(b) If  $\langle \mathbf{w}_{j,r}^{(t,b)},\pmb{\xi}_i\rangle \leq 0$ , then given that  $\langle \mathbf{w}_{j,r}^{(t,b)},\pmb{\xi}_i\rangle \approx \overline{\rho}_{j,r,i}^{(t,b)}$  and  $\overline{\rho}_{j,r,i}^{(t,b+1)} \leq \overline{\rho}_{j,r,i}^{(t,b)} + \frac{\eta(P-1)^2\|\pmb{\xi}_i\|_2^2}{Bm}$ , we can assert that, provided  $\eta$  is sufficiently small, the term  $\overline{\rho}_{j,r,i}^{(t,b)}$  can be upper bounded by a small constant.
In contrast to the analysis of SGD, which provides an upper bound for  $\overline{\rho}_{j,r,i}^{(t,b)}$  of order  $O(\log d)$ , the noise memorization prevention property described in Lemma 4.3 allows us to obtain an upper bound for  $\overline{\rho}_{j,r,i}^{(t,b)}$  of order  $O(1)$  throughout  $[0,T_1]$ . This indicates that SAM memorizes less noise compared to SGD. On the other hand, the signal coefficient  $\gamma_{j,r,i}^{(t)}$  also increases to  $\Omega (1)$  for SAM, following the same argument as in SGD. This property ensures that training with SAM does not exhibit harmful overfitting for the same signal-to-noise ratio at which training with SGD suffers from harmful overfitting.
In this section, we conduct synthetic experiments to validate our theory. Additional experiments on real data sets can be found in Appendix A.
We set training data size  $n = 20$  without label-flipping noise. Since the learning problem is rotation-invariant, without loss of generality, we set  $\pmb{\mu} = \| \pmb{\mu}\|_2 \cdot [1,0,\dots,0]^\top$ . We then generate the noise vector  $\pmb{\xi}$  from the Gaussian distribution  $\mathcal{N}(0,\sigma_p^2\mathbf{I})$  with fixed standard deviation  $\sigma_p = 1$ . We train a two-layer CNN model defined in Section 2 with the ReLU activation function. The number of filters is set as  $m = 10$ . We use the default initialization method in PyTorch to initialize the CNN parameters and train the CNN with full-batch gradient descent with a learning rate of 0.01 for 100 iterations. We consider different dimensions  $d$  ranging from 1000 to 20000, and different signal strengths  $\| \pmb{\mu}\|_2$  ranging from 0 to 10. Based on our results, for any dimension  $d$  and signal strength  $\pmb{\mu}$  setting we consider, our training setup can guarantee a training loss smaller than 0.05. After training, we estimate the test error for each case using 1000 test data points. We report the test error heat map with average results over 10 runs in Figure 2.
Sharpness Aware Minimization. Foret et al. (2020), and Zheng et al. (2021) concurrently introduced methods to enhance generalization by minimizing the loss in the worst direction, perturbed from the current parameter. Kwon et al. (2021) introduced ASAM, a variant of SAM, designed to address parameter re-scaling. Subsequently, Liu et al. (2022b) presented LookSAM, a more computationally
efficient alternative. Zhuang et al. (2022) highlighted that SAM did not consistently favor the flat minima and proposed GSAM to improve generalization by minimizing the surrogate. Recently, Zhao et al. (2022) showed that the SAM algorithm is related to the gradient regularization (GR) method when the loss is smooth and proposed an algorithm that can be viewed as a generalization of the SAM algorithm. Meng et al. (2023) further studied the mechanism of Per-Example Gradient Regularization (PEGR) on the CNN training and revealed that PEGR penalizes the variance of pattern learning.
Benign Overfitting in Neural Networks. Since the pioneering work by Bartlett et al. (2020) on benign overfitting in linear regression, there has been a surge of research studying benign overfitting in linear models, kernel methods, and neural networks. Li et al. (2021b); Montanari and Zhong (2022) examined benign overfitting in random feature or neural tangent kernel models defined in two-layer neural networks. Chatterji and Long (2022) studied the excess risk of interpolating deep linear networks trained by gradient flow. Understanding benign overfitting in neural networks beyond the linear/kernel regime is much more challenging because of the non-convexity of the problem. Recently, Frei et al. (2022) studied benign overfitting in fully-connected two-layer neural networks with smoothed leaky ReLU activation. Cao et al. (2022) provided an analysis for learning two-layer convolutional neural networks (CNNs) with polynomial ReLU activation function  $(\mathrm{ReLU}^q, q > 2)$ . Kou et al. (2023) further investigates the phenomenon of benign overfitting in learning two-layer ReLU CNNs. Kou et al. (2023) is most related to our paper. However, our work studied SGD rather than GD, which requires advanced techniques to control the update of coefficients at both batch-level and epoch-level. We also provide a novel analysis for SAM, which differs from the analysis of GD/SGD.
In this work, we rigorously analyze the training behavior of two-layer convolutional ReLU networks for both SGD and SAM. In particular, we precisely outlined the conditions under which benign overfitting can occur during SGD training, marking the first such finding for neural networks trained with mini-batch SGD. We also proved that SAM can lead to benign overfitting under circumstances that prompt harmful overfitting via SGD, which demonstrates the clear theoretical superiority of SAM over SGD. Our results provide a deeper comprehension of SAM, particularly when it comes to its utilization with non-smooth neural networks. An interesting future work is to consider other modern deep learning techniques, such as weight normalization, momentum, and weight decay, in our analysis.
We thank the anonymous reviewers for their helpful comments. ZC, JZ, YK, and QG are supported in part by the National Science Foundation CAREER Award 1906169 and IIS-2008981, and the Sloan Research Fellowship. XC and CJH are supported in part by NSF under IIS-2008173, IIS-2048280, CISCO and Sony. The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies.
AHN, K., BUBECK, S., CHEWI, S., LEE, Y. T., SUAREZ, F. and ZHANG, Y. (2022). Learning threshold neurons via the" edge of stability". arXiv preprint arXiv:2212.07469.
ALLEN-ZHU, Z. and LI, Y. (2020). Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint arXiv:2012.09816.
ALLEN-ZHU, Z. and LI, Y. (2022). Feature purification: How adversarial training performs robust deep learning. In 2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS). IEEE.
ANDRIUSHCHENKO, M. and FLAMMARION, N. (2022). Towards understanding sharpness-aware minimization. In International Conference on Machine Learning. PMLR.
ANDRIUSHCHENKO, M., VARRE, A. V., PILLAUD-VIVIEN, L. and FLAMMARION, N. (2023). Sgd with large step sizes learns sparse features. In International Conference on Machine Learning. PMLR.
MENG, X., CAO, Y. and ZOU, D. (2023). Per-example gradient regularization improves learning signals from noisy data. arXiv preprint arXiv:2303.17940.
MONTANARI, A. and ZHONG, Y. (2022). The interpolation phase transition in neural networks: Memorization and generalization under lazy training. The Annals of Statistics 50 2816-2847.
OLSHAUSEN, B. A. and FIELD, D. J. (1997). Sparse coding with an overcomplete basis set: A strategy employed by v1? Vision research 37 3311-3325.
SHEN, R., BUBECK, S. and GUNASEKAR, S. (2022). Data augmentation as feature manipulation: a story of desert cows and grass cows. arXiv preprint arXiv:2203.01572.
VERSHYNIN, R. (2018). High-Dimensional Probability: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics, Cambridge University Press.
WEN, K., MA, T. and LI, Z. (2022). How does sharpness-aware minimization minimize sharpness? arXiv preprint arXiv:2211.05729.
ZHANG, C., BENGIO, S., HARDT, M., RECHT, B. and VINYALS, O. (2021). Understanding deep learning (still) requires rethinking generalization. Communications of the ACM 64 107-115.
ZHAO, Y., ZHANG, H. and Hu, X. (2022). Penalizing gradient norm for efficiently improving generalization in deep learning. In International Conference on Machine Learning. PMLR.
ZHENG, Y., ZHANG, R. and MAO, Y. (2021). Regularizing neural networks via adversarial model perturbation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition.
ZHUANG, J., GONG, B., YUAN, L., CUI, Y., ADAM, H., DVORNEK, N., TATIKONDA, S., DUNCAN, J. and LIU, T. (2022). Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.08065.
Table 1: Top-1 accuracy  $(\%)$  of ResNet-50 on the ImageNet dataset when we vary the starting point of using the SAM update rule, baseline result is  $76.4\%$ .
τ	10%	30%	50%	70%	90%
0.01	76.9	76.9	76.9	76.7	76.7
0.02	77.1	77.0	76.9	76.8	76.6
0.05	76.2	76.4	76.3	76.3	76.2
Table 2: Top-1 accuracy (%) of wide ResNet on the CIFAR datasets when adding different levels of Gaussian noise.
Model	Noise	Dataset	Optimizer	Accuracy
WRN-16-8	-	CIFAR-10	SGD	96.69
WRN-16-8	-	CIFAR-10	SAM	97.19
WRN-16-8	N(0,0.1)	CIFAR-10	SGD	95.87
WRN-16-8	N(0,0.1)	CIFAR-10	SAM	96.57
WRN-16-8	N(0,0.3)	CIFAR-10	SGD	92.40
WRN-16-8	N(0,0.3)	CIFAR-10	SAM	93.37
WRN-16-8	N(0,1)	CIFAR-10	SGD	79.50
WRN-16-8	N(0,1)	CIFAR-10	SAM	80.37
WRN-16-8	-	CIFAR-100	SGD	81.93
WRN-16-8	-	CIFAR-100	SAM	83.68
A.1 Real Experiments on CIFAR
In this section, we provide the experiments on real data sets.