SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Assumptions. To directly compare with (Ahn et al., 2023a), we focus on the cases where  $y_0^2 - x_0^2 = \gamma / \eta$  and  $\gamma \in [\frac{1}{2}, 2]$  is a constant of moderate size; hence,  $\eta$  is not too different from the  $\eta_{\mathrm{GD}}$  defined in Theorem 3. In contrast to most prior works which assume a tiny  $\rho$  (e.g., (Wen et al., 2023)), we allow  $\rho$  to be as large as a constant (i.e., we only require  $\rho = \mathcal{O}(1)$  in Theorem 4 and Theorem 5).
3.1.1 SAM Keeps Drifting Toward the Origin
We characterize the trajectory of SAM when applied to the loss defined in Equation 5 as follows:
Theorem 4 (SAM over Single-Neuron Networks; Informal). For any  $\eta \in [\frac{1}{2}\eta_{GD}, 2\eta_{GD}]$  and  $\rho = \mathcal{O}(1)$ , the SAM trajectory over the loss function  $\mathcal{L}(x, y) = \ell(xy)$  can be divided into three phases:
1. Initial Phase.  $x_{t}$  drops so rapidly that  $|x_{t}| = \mathcal{O}(\sqrt{\eta})$  in  $\mathcal{O}(1 / \eta)$  steps. Meanwhile,  $y_{t}$  remains large: specifically,  $y_{t} = \Omega (\sqrt{1 / \eta})$ . Thus, SAM approaches the  $y$ -axis (the set of global minima).
Table 1: Limiting points of GD, SAM, and USAM for the  $\ell\left( xy\right)$  loss (assuming  $\eta  > {\eta }_{\mathrm{{GD}}}$  ).
Algorithm	GD	SAM	USAM
Limiting Point (0,y∞)	y2∞≈2/η	y2∞≈0	(1+ρy2∞)y2∞≈2/η
2. Middle Phase.  $x_{t}$  oscillates closely to the axis such that  $|x_{t}| = \mathcal{O}(\sqrt{\eta})$  always holds. Meanwhile,  $y_{t}$  decreases fast until  $y_{t} \leq |x_{t}|^{2}$  - that is,  $|x_{t}|$  remains small and SAM approaches the origin.
3. Final Phase.  $w_{t} = (x_{t},y_{t})$  gets close to the origin such that  $|x_t|,|y_t| = \mathcal{O}(\sqrt{\eta} +\eta \rho)$ . We then show that  $w_{t}$  remains in this region for the subsequent iterates.
Informally, SAM first approaches the minimizers on  $y$ -axis (which form a manifold) and then keeps moving until a specific minimum. Moreover, SAM always approaches this minimum for almost all  $(\eta, \rho)$ 's. This matches our motivating experiment in Figure 2: No matter what hyper-parameters are chosen, SAM always drift along the set of minima, in contrast to the behavior of GD. This property allows SAM always to approach the origin  $(0,0)$  and remains in its neighborhood, while GD converges to  $(0, \sqrt{2/\eta})$  (see Table 1). The formal version of Theorem 4 is in Appendix F.
3.1.2 USAM Gets Trapped at Different Minima
We move on to characterize the dynamics of USAM near the minima. Like GD or SAM, the first few iterations of USAM drive iterates to the  $y$ -axis. However, unlike SAM, USAM does not keep drifting along the  $y$ -axis and stops at some threshold – in the result below, we prove a lower bound on  $y_{t}^{2}$  that depends on both  $\eta$  and  $\rho$ . In other words, the lack of normalization factor leads to diminishing drift.
Theorem 5 (USAM over Single-Neuron Networks; Informal). For any  $\eta \in [\frac{1}{2}\eta_{GD}, 2\eta_{GD}]$  and  $\rho = \mathcal{O}(1)$ , the USAM trajectory over the loss function  $\mathcal{L}(x, y) = \ell(xy)$  have the following properties:
1. Initial Phase. Similar to Initial Phase of Theorem 4,  $|x_{t}| = \mathcal{O}(\sqrt{\eta})$  and  $y_{t} = \Omega (\sqrt{1 / \eta})$  hold for the first  $\mathcal{O}(1 / \eta)$  steps. That is, USAM also approaches  $y$ -axis, the set of global minima.
2. Final Phase. However, for USAM, once the following condition holds for some round  $t$ :<sup>3</sup>
$$
\left. \left(1 + \rho y _ {\mathrm {t}} ^ {2}\right) y _ {\mathrm {t}} ^ {2} <   \frac {2}{\eta}, \quad i . e ., \quad y _ {\mathrm {t}} ^ {2} <   \tilde {y} _ {U S A M} ^ {2} \triangleq \left(\sqrt {8 \frac {\rho}{\eta} + 1} - 1\right) / 2 \rho , \right. \tag {7}
$$
$\left|x_{t}\right|$  decays exponentially fast, which in turn ensures  $y_{\infty}^{2} \triangleq \lim \inf_{t \to \infty} y_{t}^{2} \gtrsim (1 - \eta^{2} - \rho^{2})y_{t}^{2}$ .
Remark. Note that USAM becomes GD as we send  $\rho \to 0+$ , and our characterized threshold  $y_{\mathrm{USAM}}^2$  indeed recovers that of GD (i.e.,  $2/\eta$  from Theorem 3) because  $\lim_{\rho \to 0+} (\sqrt{8\rho/\eta + 1} - 1)/2\rho = 2/\eta$ .
Compared with SAM, the main difference occurs when close to minima, i.e.,  $|x_{t}| = \mathcal{O}(\sqrt{\eta})$ . Consistent with our motivating experiment (Figure 2), the removal of normalization leads to diminishing drift along the minima. Thus, USAM is more like an improved version of GD rather than a simplification of SAM, and the comparison between Theorem 3 and Theorem 5 reveals that USAM only improves over GD if  $\rho$  is large enough – in which case USAM is prone to diverges as we discussed in Section 2.
See Appendix G for a formal version of Theorem 5 together with its proof.
3.1.3 Technical Distinctions Between GD, SAM, and USAM
Before moving on, we present a more technical comparison between the results stated in Theorem 3 versus Theorem 4 and Theorem 5. We start with an intuitive explanation of why GD and USAM get stuck near the manifold of minima but SAM does not: when the iterates approach the set of minima, both  $w_{t}$  and  $\nabla \mathcal{L}(w_t)$  become small. Hence the normalization plays an important role: as  $\nabla \mathcal{L}(w_t)$  are small,  $w_{t}$  and  $w_{t} + \rho \nabla \mathcal{L}(w_{t})$  become nearly identical, which leads to a diminishing updates of GD and USAM near the minima. On the other hand, having the normalization term, the SAM update doesn't diminish, which prevents SAM from converging to a minimum and keeps drifting along the manifold.
${}^{2}$  Specifically,either an exponential decay  ${y}_{t + 1} \lesssim  {\left( 1 - \eta {\rho }^{2}\right) }{y}_{t}$  or a constant drop  ${y}_{t + 1} \lesssim  {y}_{t} - \eta \rho$  appears  ${}^{3}$  The  ${\widetilde{y}}_{\text{USAM }}$  in Equation 7 is defined as the solution to the equation  ${\left( 1 + \rho {y}^{2}\right) }{y}^{2} = 2$  .
This high-level idea is supported by the following calculation: recall Equation 6 that  $\nabla \mathcal{L}(x_t, y_t) = \ell'(x_t y_t) \cdot [y_t x_t]^\top$ . Hence, when  $|x_t| \ll y_t$  in Final Phase, the "ascent gradient" direction  $\nabla \mathcal{L}(x_t, y_t)$  (i.e., the ascent steps in Equation 1 and Equation 2) is almost perpendicular to the  $y$ -axis. We thus rewrite the update direction (i.e., the difference between  $w_{t+1}$  and  $w_t$ ) for each algorithm as follows.
1. For SAM, after normalization,  $\frac{\nabla\mathcal{L}(w_t)}{\|\nabla\mathcal{L}(w_t)\|}$  is roughly a unit vector along the  $x$ -axis. Hence, the update direction is the gradient at  $w_{t + 1 / 2} \approx [\rho y_t]^\top$ . Once  $y_t$  is large (making  $w_{t + 1 / 2}$  far from minima),  $\nabla \mathcal{L}(w_{t + 1 / 2})$  thus have a large component along  $y_t$ , which leads to drifting near minima.
2. For GD, by approximating  $\ell'(u) \approx u$ , we derive  $\nabla \mathcal{L}(x_t, y_t) \approx [x_t y_t^2 - x_t^2 y_t]^\top$ . When  $2/\eta > y_t^2$ , the magnitude of  $x_t$  is updated as  $|x_{t+1}| \approx |x_t - \eta x_t y_t^2| = (1 - \eta y_t^2)x_t$ , which allows an exponential decay. Thus, GD converges to a minimum and stop moving soon after  $2/\eta > y_t^2$ .
3. For USAM, the descent gradient is taken at  $w_{t} + \rho \nabla \mathcal{L}(w_{t}) \approx [(1 + \rho y_{t}^{2})x_{t} - (1 + \rho x_{t}^{2})y_{t}]^{\top}$ . Thus,  $\nabla \mathcal{L}(w_{t} + \rho \nabla \mathcal{L}(w_{t})) \approx [(1 + \rho y_{t}^{2})(1 + \rho x_{t}^{2})^{2}x_{t}y_{t}^{2} - (1 + \rho y_{t}^{2})^{2}(1 + \rho x_{t}^{2})x_{t}^{2}y_{t}]^{\top}$  by writing  $\ell'(u) \approx u$ . This makes USAM deviate away from SAM and behave like GD: by the similar argument as GD, USAM stops at a minimum soon after  $2/\eta > (1 + \rho y_{t}^{2})(1 + \rho x_{t}^{2})^{2}y_{t}^{2} \approx (1 + \rho y_{t}^{2})y_{t}^{2}!$
Hence, the normalization factor in the ascent gradient helps maintain a non-diminishing component along the minima, leading SAM to keep drifting. This distinguishes SAM from GD and USAM.
3.2 USAM Gets Trapped Once Close to Minima
In this section, we extend our arguments to nonconvex costs satisfying Polyak-Lojasiewicz (PL) functions (see, e.g., (Karimi et al., 2016)). Recall that  $f$  satisfies the  $\mu$ -PL condition if  $\frac{1}{2} \| \nabla \mathcal{L}(w) \|^2 \geq \mu(\mathcal{L}(w) - \min_w \mathcal{L}(w))$  for all  $w$ . Building upon the analysis of Andriushchenko and Flammarion (2022), we show the following result when applying USAM to  $\beta$ -smooth and  $\mu$ -PL losses.
Theorem 6 (USAM over PL Losses; Informal). For  $\beta$ -smooth and  $\mu$ -PL loss  $\mathcal{L}$ , for any  $\eta < 1/\beta$  and  $\rho < 1/\beta$ , and for any initialization  $w_0$ ,  $\| w_t - w_0 \| \leq \text{poly}(\eta, \rho, \beta, \mu) \cdot \sqrt{\mathcal{L}(w_0) - \min_w \mathcal{L}(w)}$ ,  $\forall t$ .
This theorem has the following consequence: Suppose that USAM encounters a point  $w_0$  that is close to some minimum (i.e.,  $\mathcal{L}(w_0) \approx \min_w \mathcal{L}(w)$ ) during training. Then Theorem 6 implies that the total distance traveled by USAM from  $w_0$  is bounded – in other words, the distance USAM moves along the manifold of minimizers can only be of order  $\mathcal{O}(\sqrt{\mathcal{L}(w_0) - \min_w \mathcal{L}(w)})$ .
As a remark, we compare Theorem 6 with the recent result by Wen et al. (2023): their result essentially implies that, for small enough  $\eta$  and  $\rho$ , SAM iterates initialized close to a manifold of the minimizers approximately track some continuous dynamics (more precisely, a Riemannian gradient flow induced by a "sharpness" measure they find) and keep drifting along the manifold. This property is indeed in sharp contrast with USAM whose total travel distance is bounded.
The formal statement and proof of Theorem 6 are contained in Appendix H.
3.3 Experiments for Practical Neural Networking Training
Figure 5: Training ResNet18 on CIFAR-10 from a bad global minimum ( $\eta = 0.001$ , batch size = 128).
We close this section by verifying our claims in practical neural network training. We train a ResNet18 on the CIFAR-10 dataset, initialized from a poor global minimum generated as per Liu et al. (2020)
(we used the "adversarial checkpoint" released by Damian et al. (2021)). This initialization has  $100\%$  training accuracy but only  $48\%$  test accuracy - which lets us observe a more pronounced algorithmic behavior near the minima via tracking the test accuracy. From Figure 5, we observe:
1. GD gets stuck at this adversarial minimum, in the sense that the test accuracy stays at  $48\%$ .
2. SAM keeps drifting while staying close to the manifold of minimizers (because the training accuracy remains  $100\%$ ), which results in better solutions (i.e., the test accuracy keeps increasing).
3. USAM with small  $\rho$  gets stuck like GD, while USAM with larger  $\rho$ 's deviate from this manifold.
Hence, USAM faces the dilemma that we describe in Subsection 3.1: a conservative hyper-parameter configuration does not lead to much drift along the minima, while a more aggressive choice easily leads to divergence. However, the stability of SAM is quite robust to the choice of hyper-parameter and they all seem to lead to consistent drift along the minima.
Remark. Apart from the "adversarial checkpoint" which is unrealistic but can help highlight different algorithms' behavior when they are close to a bad minimum, we also conduct the same experiments but instead initialized from a "full-batch checkpoint" (Damian et al., 2021), which is the  $100\%$  training accuracy point reached by running full-batch GD on the training loss function. The result is plotted as Figure 8 in Subsection B.1. One can observe that USAM still gets stuck at the "full-batch checkpoint", while SAM keeps increasing its test accuracy via drifting along the minima manifold.
4 Case Study: Learning Threshold Neurons for Sparse Coding Problem
To incorporate our two findings into a single example, we consider training one-hidden-layer ReLU networks for the sparse coding problem, a setup considered in (Ahn et al., 2023a) to study the role of  $\eta$  in GD. Without going into details, the crux of their experiment is to understand how GD with large  $\eta$  finds desired structures of the network – in this specific case, the desired structure is the negative bias in ReLU unit (also widely known as “thresholding unit/neuron”). In this section, we evaluate SAM and USAM under the same setup, illustrating the importance of normalization.
Main observation of Ahn et al. (2023a). Given this background, the main observation of Ahn et al. (2023a) is that i) when training the ReLU network with GD, different learning rates induce very different trajectories; and ii) the desired structure, namely a negative bias in ReLU, only arises with large "unstable" learning rates for which GD exhibits unstable behaviors. We reproduce their results in Figure 6, plotting the test accuracy on the left and the bias of ReLU unit on the right. As they claimed, GD with larger  $\eta$  learns more negative bias, which leads to better test accuracy.
Figure 6: Behavior of GD for sparse coding problem.
Their inspiring observation is however a bit discouraging for practitioners. According to their theoretical results, such learning rates have to be quite large – large to the point where GD shows very unstable behavior (à la Edge-of-Stability (Cohen et al., 2021)). In practice, without knowledge of the problem, this requires a careful hyper-parameter search to figure out the correct learning rate. More importantly, such large and unstable learning rates may cause GD to diverge or lead to worse performance. More discussions can be found in the recent paper by Kaur et al. (2022).
In contrast, as we will justify shortly, SAM does not suffer from such a "dilemma of tuning" – matching with our results in Theorem 4. Moreover, the removal of normalization no longer attains such a property, as we demonstrated in Theorem 5. In particular, for USAM, one also needs to carefully tune  $\eta$  and  $\rho$  for better performance – as we inspired in Theorem 5 and Theorem 6, small  $(\eta, \rho)$  makes the iterates get stuck early; on the other hand, as we presented in Section 2, an aggressive choice causes USAM to diverge. The following experiments illustrate these claims in more detail.
In Figure 7, we plot the performance of SAM, USAM, and GD with different  $\eta$ 's (while fixing  $\rho$ ) — gray lines for GD, solid lines for SAM, and dashed lines for USAM. From the plot, USAM behaves more similarly to GD than SAM: the bias does not decrease sufficiently when the learning rate is not large enough, which consequently leads poor test accuracy. On the other hand, no matter what
Figure 7: Behaviors of different algorithms for sparse coding problem  $(\rho = 0.01)$ . The gray curves (corresponding to GD with different learning rates) are taken from Figure 6 with the same set of  $\eta$ 's.
$\eta$  is chosen for SAM, bias is negative enough and ensures better generalization. Hence, Figure 7 illustrates that compared to SAM, USAM is less robust to the tuning of  $\eta$ .
In Figure 9 (deferred to Subsection B.2), we also compare these three algorithms when varying  $\rho$  and fixing  $\eta$ . In addition to what we observe in Figure 7, we show that normalization also helps stability - USAM quickly diverges as we increase  $\rho$ , while SAM remains robust to the choice of  $\rho$ . Thus, USAM is also less robust to the tuning of  $\rho$ . In other words, our observation in Figure 7 extends to  $\rho$ .
Hence, putting Figure 7 and Figure 9 together, we conclude that SAM is robust to different configurations of  $(\eta, \rho)$  while USAM is robust to neither of them. Hence, the normalization of SAM eases hyper-parameter tuning, which is typically a tough problem for GD and many other algorithms - normalization boosts the success of SAM in practice.
In this paper, we investigate the role played by normalization in SAM. By theoretically characterizing the behavior of SAM and USAM on both convex and non-convex losses and empirically verifying our conclusions via real-world neural networks, we found that normalization i) helps stabilize the algorithm iterates, and ii) enables the algorithm to keep moving along the manifold of minimizers, leading to better performance in many cases. Moreover, as we demonstrate via various experiments, these two properties make SAM require less hyper-parameter tuning, supporting its practicality.
In this work, we follow a recent research paradigm of "physics-style" approaches to understanding deep neural networks based on simplified models (c.f. (Zhang et al., 2022; Garg et al., 2022; von Oswald et al., 2023; Abernethy et al., 2023; Allen-Zhu and Li, 2023; Liu et al., 2023; Li et al., 2023; Ahn et al., 2023a,b,c)). We found such physics-style approaches quite helpful, especially for complex modern neural networks. We hope that our work builds stepping stones for future works on understanding working mechanisms of modern deep neural networks.
Kwangjun Ahn was supported by the ONR grant (N00014-20-1-2394) and MIT-IBM Watson as well as a Vannevar Bush fellowship from Office of the Secretary of Defense. Suvrit Sra acknowledges support from an NSF CAREER grant (1846088), and NSF CCF-2112665 (TILOS AI Research Institute). Kwangjun Ahn also acknowledges support from the Kwanjeong Educational Foundation. We thank Kaiyue Wen, Zhiyuan Li, and Hadi Daneshmand for their insightful discussions.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generalization in deep learning. Advances in neural information processing systems, 30, 2017.
Dongkuk Si and Chulhee Yun. Practical sharpness-aware minimization cannot converge all the way to optima. NeurIPS 2023 (arXiv:2306.09850), 2023.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vlademyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174. PMLR, 2023.
Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along gd trajectory: Progressive sharpening and edge of stability. Advances in Neural Information Processing Systems, 35:9983-9994, 2022.
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? In International Conference on Learning Representations, 2023.
Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization. Advances in Neural Information Processing Systems, 33:2958-2969, 2020.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301, 2022.
Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8156-8165, 2021.
Qihuang Zhong, Liang Ding, Li Shen, Peng Mi, Juhua Liu, Bo Du, and Dacheng Tao. Improving sharpness-aware minimization with fisher mask for better generalization on language models. In Findings of the Association for Computational Linguistics: EMNLP 2022, pages 4064-4085, 2022.
Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha C Dvornek, James Duncan, Ting Liu, et al. Surrogate gap minimization improves sharpness-aware training. In International Conference on Learning Representations, 2022.
A Setup of the Motivating Experiment 15
B Additional Experimental Results 15
B.1 Running SAM and USAM from Other Initializations 15
B.2 Varying  $\rho$  While Fixing  $\eta$  in Sparse Coding Example 15
C Omitted Proof for Smooth and Strongly Convex Losses 16
C.1 SAM Allows a Descent Lemma Like GD 16
C.2 USAM Diverges on Quadratic Losses 17
D Omitted Proof for Scalar Factorization Problems 18
D.1 SAM Always Converges on Scalar Factorization Problems 18
D.2 USAM Diverges with Small  $\rho$  19
E Assumptions in the Single-Neuron Linear Network Model 20
F Omitted Proof of SAM Over Single-Neuron Linear Networks 20
F.1 Basic Properties and Notations 20
F.2 Initial Phase:  $x_{t}$  Decreases Fast while  $y_{t}$  Remains Large 21
F.3 Middle Phase:  $y$  Keeps Decreasing Until Smaller Than  $|x_{t}|$  23
F.4 Final Phase: Both  $x_{t}$  and  $y_{t}$  Oscillates Around the Origin 24
G Omitted Proof of USAM Over Single-Neuron Linear Networks 25
G.1 Basic Properties and Notations 25
G.2 Initial Phase:  $x_{t}$  Decreases Fast while  $y_{t}$  Remains Large 26
G.3 Final Phase:  $y_{t}$  Gets Trapped Above the Origin 28
H Omitted Proof of USAM Over General PL functions 29
A Setup of the Motivating Experiment
In the motivating experiments (Figure 1 and Figure 2), we follow the over-parameterized matrix sensing setup as Li et al. (2018) and Blanc et al. (2020). Specifically, we do the following:
1. Generate the true matrix by sampling each entry of  $U^{\star} \in \mathbb{R}^{d \times r}$  independently from a standard Gaussian distribution and let  $X^{\star} = \bar{U}^{\star}(U^{\star})^{\top}$ .
2. Normalize each column of  $U^{\star}$  to unit norm so that the spectral norm of  $U^{\star}$  is close to one.
3. For every sensing matrix  $A_{i}$  ( $i = 1,2,\dots,m$ ), sample the entries of  $A_{i}$  independently from a standard Gaussian distribution. Then observe  $b_{i} = \langle A_{i},X^{\star}\rangle$ .
In particular, for the experiments, we chose  $r = 5$ ,  $d = 100$ , and  $m = 5dr$ .
B.1 Running SAM and USAM from Other Initializations
Figure 8: Training ResNet18 on CIFAR-10 from a "full-batch checkpoint" of Damian et al. (2021) ( $\eta = 0.001$ , batch size = 128).
B.2 Varying  $\rho$  While Fixing  $\eta$  in Sparse Coding Example
Figure 9: Behaviors of different algorithms for sparse coding problem ( $\eta = 0.025$ ). Note that USAM with  $\rho = 0.5$  (the dashed violet curve) diverges and becomes invisible except for the very beginning.
Aside from Figure 7 which varies  $\eta$  while fixing  $\rho$ , we perform the same experiment when varying  $\rho$  and fixing  $\eta$ . The main observation for SAM is similar to that of Figure 7: different hyper-parameters all keep decreasing the bias and give better test accuracy - even with the tiniest choice  $\rho = 0.01$ .
However, for USAM, there are three different types of  $\rho$ 's as shown in Figure 9:
1. For tiny  $\rho = 0.01$ , the bias doesn't decrease much. Consequently, the performance of USAM nearly degenerates to that of GD - while SAM with  $\rho = 0.01$  still gives outstanding performance.
2. For moderate  $\rho = 0.1$  and  $\rho = 0.3$ , USAM manages to decrease the bias and improve its accuracy, though with a slower speed than SAM with the same choices of  $\rho$ .
3. For large  $\rho = 0.5$  (where SAM still works well; see the solid curve in blue), USAM diverges.
Thus, the dilemma described in Section 4 indeed applies to not only  $\eta$  but also  $\rho$  - matching our main conclusion that normalization helps make hyper-parameter tuning much more manageable.
C Omitted Proof for Smooth and Strongly Convex Losses
We shall first restate Theorem 1 here for the ease of presentation.
Theorem 7 (Restatement of Theorem 1). For any  $\alpha$ -strongly-convex and  $\beta$ -smooth loss function  $\mathcal{L}$ , for any learning rate  $\eta \leq 2 / \beta$  and perturbation radius  $\rho \geq 0$ , the following holds:
$$
\mathcal {L} (w _ {t}) - \mathcal {L} (w ^ {\star}) \leq \left(1 - \alpha \eta (2 - \eta \beta)\right) ^ {t} (\mathcal {L} (w _ {0}) - \mathcal {L} (w ^ {\star})) + \frac {\eta \beta^ {3} \rho^ {2}}{2 \alpha (2 - \eta \beta)}, \quad \forall t.
$$
We will show these two conclusions separately. The proof directly follows from Theorem 8 and Theorem 10.
C.1 SAM Allows a Descent Lemma Like GD
Theorem 8 (SAM over Strongly Convex and Smooth Losses). For any  $\alpha$ -strongly-convex and  $\beta$ -smooth loss function  $\mathcal{L}$ , for any learning rate  $\eta \leq 2 / \beta$  and perturbation radius  $\rho \geq 0$ , the following holds for SAM: