SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yan Dai *
IIIS, Tsinghua University
yan-dai20@mails.tsinghua.edu.cn
Kwangjun Ahn *
EECS, MIT
Suvrit Sra
TU Munich / MIT
suvrit@mit.edu
Sharpness-Aware Minimization (SAM) is a recently proposed gradient-based optimizer (Foret et al., ICLR 2021) that greatly improves the prediction performance of deep neural networks. Consequently, there has been a surge of interest in explaining its empirical success. We focus, in particular, on understanding the role played by normalization, a key component of the SAM updates. We theoretically and empirically study the effect of normalization in SAM for both convex and non-convex functions, revealing two key roles played by normalization: i) it helps in stabilizing the algorithm; and ii) it enables the algorithm to drift along a continuum (manifold) of minima – a property identified by recent theoretical works that is the key to better performance. We further argue that these two properties of normalization make SAM robust against the choice of hyper-parameters, supporting the practicality of SAM. Our conclusions are backed by various experiments.
We study the recently proposed gradient-based optimization algorithm Sharpness-Aware Minimization (SAM) (Foret et al., 2021) that has shown impressive performance in training deep neural networks to generalize well (Foret et al., 2021; Bahri et al., 2022; Mi et al., 2022; Zhong et al., 2022). SAM updates involve an ostensibly small but key modification to Gradient Descent (GD). Specifically, for a loss function  $\mathcal{L}$  and each iteration  $t \geq 0$ , instead of updating the parameter  $w_{t}$  as  $w_{t+1} = w_{t} - \eta \nabla \mathcal{L}(w_{t})$  (where  $\eta$  is called the learning rate), SAM performs the following update:
$$
w _ {t + 1} = w _ {t} - \eta \nabla \mathcal {L} \left(w _ {t} + \rho \frac {\nabla \mathcal {L} \left(w _ {t}\right)}{\| \nabla \mathcal {L} \left(w _ {t}\right) \|}\right), \tag {1}
$$
where  $\rho$  is an additional hyper-parameter that we call the perturbation radius. Foret et al. (2021) motivate SAM as an algorithm minimizing the robust loss  $\max_{\| \epsilon \| \leq \rho}\mathcal{L}(w + \epsilon)$ , which is roughly the loss at  $w$  (i.e.,  $\mathcal{L}(w)$ ) plus the "sharpness" of the loss landscape around  $w$ , hence its name.
The empirical success of SAM has driven a recent surge of interest in characterizing its dynamics and theoretical properties (Bartlett et al., 2022; Wen et al., 2023; Ahn et al., 2023d). However, a major component of SAM remains unexplained in prior work: the role and impact of the normalization factor  $\frac{1}{\|\nabla\mathcal{L}(w_t)\|}$  used by SAM. In fact, quite a few recent works drop the normalization factor for simplicity when analyzing SAM (Andriushchenko and Flammarion, 2022; Behdin and Mazumder, 2023; Agarwala and Dauphin, 2023; Kim et al., 2023; Compagnoni et al., 2023). Instead of the SAM update (1), these works consider the following update that we call Un-normalized SAM (USAM):
$$
w _ {t + 1} = w _ {t} - \eta \nabla \mathcal {L} \left(w _ {t} + \rho \nabla \mathcal {L} \left(w _ {t}\right)\right). \tag {2}
$$
*The first two authors contribute equally. Work done while Yan Dai was visiting MIT.
In principle, the normalization in Equation 1 may make SAM ill-defined. However, Wen et al. (2023, Appendix B) showed that except for countably many learning rates, SAM (with any  $\rho$ ) is always well-defined for almost all initialization. Hence, throughout the paper, we assume that the SAM iterates are always well-defined.
Apart from experimental justifications in (Andriushchenko and Flammarion, 2022), the effect of this simplification has not yet been carefully investigated, although it is already widely adopted in the community. Thus, is it really the case that such normalization can be omitted "for simplification" when theoretically analyzing SAM? These observations raise our main question:
$$
\text {W h i t} \frac {1}{\| \nabla \mathcal {L} (w _ {t}) \|} \text {i n} (1)?
$$
1.1 Motivating Experiments and Our Contributions
We present our main findings through two motivating experiments. For the setting, we choose the well-known over-parameterized matrix sensing problem (Li et al., 2018); see Appendix A for details.
1. Normalization helps with stability. We first pick a learning rate  $\eta$  that allows GD to converge, and we gradually increase  $\rho$  from 0.001 to 0.1. Considering the early stage of training shown in Figure 1. One finds that SAM has very similar behavior to GD, whereas USAM diverges even with a small  $\rho$ -it seems that normalization helps stabilize the algorithm.
Figure 1: Role of normalization for stabilizing algorithms  $(\eta = 0.05)$ .
2. Normalization permits moving along minima. We reduce the step size by 10 times and consider their performance of reducing test losses in the long run. One may regard Figure 2 as the behavior of SAM, USAM, and GD when close to a "manifold" of minima (which exists since the problem is over-parametrized) as the training losses are close to zero. The first plot compares SAM and USAM with the same  $\rho = 0.1$  (the largest  $\rho$  for which USAM doesn't diverge): notice that USAM and GD both converge to a minimum and do not move further; on the other hand, SAM keeps decreasing the test loss, showing its ability to drift along the manifold. We also vary  $\rho$  and compare their behaviors (shown on the right): the ability of SAM to travel along the manifold of minimizers seems to be robust to the size of  $\rho$ , while USAM easily gets stuck at a minimum.
Figure 2: Role of normalization when close to a manifold of minimizers  $(\eta = 0.005)$ .
Overview of our contributions. In this work, as motivated by Figure 1 and Figure 2, we identify and theoretically explain the two roles of normalization in SAM. The paper is organized as follows.
1. In Section 2, we study the role of normalization in the algorithm's stability and show that normalization helps stabilize. In particular, we demonstrate that normalization ensures that GD's convergence implies SAM's non-divergence, whereas USAM can start diverging much earlier.
2. In Section 3, we study the role of normalization near a manifold of minimizers and show that the normalization factor allows iterates to keep drifting along this manifold - giving better performance in many cases. Without normalization, the algorithm easily gets stuck and no longer progresses.
3. In Section 4, to illustrate our main findings, we adopt the sparse coding example of Ahn et al. (2023a). Their result implies a dilemma in hyper-parameter tuning for GD: a small  $\eta$  gives worse performance, but a large  $\eta$  results in divergence. We show that this dilemma extends to USAM - but not SAM. In other words, SAM easily solves the problem where GD and USAM often fail.
These findings also shed new light on why SAM is practical and successful, as we highlight below.
Practical importance of our results. The main findings in this work explain and underscore several practical aspects of SAM that are mainly due to the normalization step. One practical feature of SAM is the way the hyper-parameter  $\rho$  is tuned: Foret et al. (2021) suggest that  $\rho$  can be tuned independently after tuning the parameters of base optimizers (including learning rate  $\eta$ , momentum  $\beta$ , and so on). In particular, this feature makes SAM a perfect "add-on" to existing gradient-based optimizers. Our findings precisely support this practical aspect of SAM. Our results suggest that the stability of SAM is less sensitive to the choice of  $\rho$ , thanks to the normalization factor.
The same principle applies to the behavior of the algorithm near the minima: Recent theoretical works (Bartlett et al., 2022; Wen et al., 2023; Ahn et al., 2023d) have shown that the drift along the manifold of minimizers is a main feature that enables SAM to reduce the sharpness of the solution (which is believed to boost generalization ability in practice) – our results indicate that the ability of SAM to keep drifting along the manifold is independent of the choice of  $\rho$ , again owing to normalization. Hence, our work suggests that the normalization factor plays an important role towards SAM's empirical success.
Sharpness-Aware Optimizers. Inspired by the empirical and theoretical observation that the generalization effect of a deep neural network is correlated with the "sharpness" of the loss landscape (see (Keskar et al., 2017; Jastrzkebski et al., 2017; Jiang et al., 2020) for empirical observations and (Dziugaite and Roy, 2017; Neyshabur et al., 2017) for theoretical justifications), several recent papers (Foret et al., 2021; Zheng et al., 2021; Wu et al., 2020) propose optimizers that penalize the sharpness for the sake of better generalization. Subsequent efforts were made on making such optimizers scale-invariant (Kwon et al., 2021), more efficient (Liu et al., 2022; Du et al., 2022), and generalize better (Zhuang et al., 2022). This paper focuses on the vanilla version proposed by Foret et al. (2021).
Theoretical Advances on SAM. Despite the success of SAM in practice, theoretical understanding of SAM was absent until two recent works: Bartlett et al. (2022) analyze SAM on locally quadratic losses and identify a component reducing the sharpness  $\lambda_{\mathrm{max}}(\nabla^2\mathcal{L}(w_t))$ , while Wen et al. (2023) characterize SAM near the manifold  $\Gamma$  of minimizers and show that SAM follows a Riemannian gradient flow reducing  $\lambda_{\mathrm{max}}(\nabla^2\mathcal{L}(w))$  when i) initialized near  $\Gamma$ , and ii)  $\eta$  is "small enough". Note that while the results of Wen et al. (2023) apply to more general loss functions, our result in Theorem 16 applies to i) any initialization far from the origin, and ii) any  $\eta = o(1)$  and  $\rho = O(1)$ . A recent work by Ahn et al. (2023d) formulates the notion of  $\varepsilon$ -approximate flat minima and analyzed the iteration complexity of practical algorithms like SAM to find such approximate flat minima. A concurrent work by Si and Yun (2023) also analyzes the original version of SAM with the normalization in (1), and makes a case that practical SAM does not converge all the way to optima.
Unnormalized SAM (USAM). USAM was first proposed by Andriushchenko and Flammarion (2022) who observed a similar performance between USAM and SAM when training ResNet over CIFAR-10. This simplification is further accepted by Behdin and Mazumder (2023) who study the regularization effect of USAM over a linear regression model, by Agarwala and Dauphin (2023) who study the initial and final dynamics of USAM over a quadratic regression model, and by Kim et al. (2023) who study the convergence instability of USAM near saddle points. To our knowledge, (Compagnoni et al., 2023) is the only work comparing SAM and USAM dynamics. More precisely, they consider the continuous-time behavior of SGD, SAM, and USAM and find different behaviors of SAM and USAM: USAM attracts local minima while SAM aims at global ones. Still, we remark that as they are considering continuous-time variants of algorithms while we consider discrete (original) versions, our results directly apply to the SAM deployed in practice and the USAM studied in theory.
Edge-of-Stability. In the optimization theory literature, Gradient Descent (GD) was only shown to find minima if the learning rate  $\eta$  is smaller than an "Edge-of-Stability" threshold, which is related to the sharpness of the nearest minimum. However, people recently observe that when training neural networks, GD with a  $\eta$  much larger than that threshold often finds good minima as well (see (Cohen
et al., 2021) and references therein). Aside from convergence, GD with large  $\eta$  is also shown to find flatter minima (Arora et al., 2022; Ahn et al., 2022; Wang et al., 2022; Damian et al., 2023).
2 Role of Normalization for Stability
In this section, we discuss the role of normalization in the stability of the algorithm. We begin by recalling a well-known fact about the stability of GD: for a convex quadratic cost with the largest eigenvalue of Hessian being  $\beta$  (i.e.,  $\beta$ -smooth), GD converges to a minimum iff  $\eta < 2/\beta$ . Given this folklore fact, we ask: how do the ascent steps in SAM (1) and USAM (2) affect their stability?
2.1 Strongly Convex and Smooth Losses
Consider an  $\alpha$ -strongly-convex and  $\beta$ -smooth loss function  $\mathcal{L}$  where GD is guaranteed to converge once  $\eta < 2 / \beta$ . We characterize the stability of SAM and USAM in the following result.
Theorem 1 (Strongly Convex and Smooth Losses). For any  $\alpha$ -strongly-convex and  $\beta$ -smooth loss function  $\mathcal{L}$ , for any learning rate  $\eta < 2/\beta$  and perturbation radius  $\rho \geq 0$ , the following holds:
1. SAM. The iterate  $w_{t}$  converges to a local neighborhood around the minimizer  $w^{\star}$ . Formally,
$$
\mathcal {L} \left(w _ {t}\right) - \mathcal {L} \left(w ^ {\star}\right) \leq \left(1 - \alpha \eta (2 - \eta \beta)\right) ^ {t} \left(\mathcal {L} \left(w _ {0}\right) - \mathcal {L} \left(w ^ {\star}\right)\right) + \frac {\eta \beta^ {3} \rho^ {2}}{2 \alpha (2 - \eta \beta)}, \quad \forall t. \tag {3}
$$
2. USAM. In contrast, there exists some  $\alpha$ -strongly-convex and  $\beta$ -smooth loss  $\mathcal{L}$  such that the USAM with  $\eta \in (2 / (\beta + \rho \beta^2), 2 / \beta]$  diverges for all except measure zero initialization  $w_0$ .
As we discussed, it is well-known that GD converges iff  $\eta < 2 / \beta$ , and Theorem 1 shows that SAM also does not diverge and stays within an  $\mathcal{O}(\sqrt{\eta}\rho)$ -neighborhood around the minimum as long as  $\eta < 2 / \beta$ . However, USAM diverges with an even lower learning rate:  $\eta > 2 / (\beta + \rho \beta^2)$  can already make USAM diverge. Intuitively, the larger the value of  $\rho$ , the easier it is for USAM to diverge.
One may notice that Equation 3, compared to the standard convergence rate of GD, exhibits an additive bias term of order  $\mathcal{O}(\eta \rho^2)$ . This term arises from the unstable nature of SAM: the perturbation in (1) (which always has norm  $\rho$ ) prevents SAM from decreasing the loss monotonically. Thus, SAM can only approach a minimum up to a neighborhood. For this reason, in this paper whenever we say SAM "finds" a minimum, we mean its iterates approach and stay within a neighborhood of that minimum.
Due to space limitations, the full proof is postponed to Appendix C and we only outline it here.
Proof Sketch. For SAM, we show an analog to the descent lemma of GD as follows (see Lemma 9):
$$
\mathcal {L} \left(w _ {t + 1}\right) \leq \mathcal {L} \left(w _ {t}\right) - \frac {1}{2} \eta (2 - \eta \beta) \| \nabla \mathcal {L} \left(w _ {t}\right) \| ^ {2} + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2}. \tag {4}
$$
By invoking the strong convexity that gives  $\mathcal{L}(w_t) - \mathcal{L}(w^*) \leq \frac{1}{2\alpha} \| \nabla f(w_t) \|^2$ , we obtain
$$
\mathcal {L} \left(w _ {t + 1}\right) - \mathcal {L} \left(w ^ {\star}\right) \leq \left(1 - \alpha \eta (2 - \eta \beta)\right) \left(\mathcal {L} \left(w _ {t}\right) - \mathcal {L} \left(w ^ {\star}\right)\right) + \frac {\eta^ {2} \beta^ {3} \rho^ {2}}{2}.
$$
Recursively applying this relation gives the first conclusion. For USAM, we consider the quadratic loss function same as (Bartlett et al., 2022). Formally, suppose that  $\mathcal{L}(w) = \frac{1}{2} w^\top \Lambda w$  where  $\Lambda = \mathrm{diag}(\lambda_1, \lambda_2, \ldots, \lambda_d)$  is a PSD matrix such that  $\lambda_1 > \lambda_2 \geq \dots \geq \lambda_d > 0$ . Let the eigenvectors corresponding to  $\lambda_1, \lambda_2, \ldots, \lambda_d$  be  $e_1, e_2, \ldots, e_d$ , respectively. Then we show the following in Theorem 10: for any  $\eta(\lambda_1 + \rho \lambda_1^2) > 2$  and  $\langle w_0, e_1 \rangle \neq 0$ , USAM must diverge. As  $\mathcal{L}(w) = \frac{1}{2} w^\top \Lambda w$  is  $\lambda_1$ -smooth and  $\lambda_d$ -strongly-convex, the second conclusion also follows.
Intuitively, the difference in stability can be interpreted as follows: during the early stage of training,  $w_{t}$  and  $\nabla \mathcal{L}(w_t)$  often have large norms. The normalization in SAM then makes the ascent step  $w_{t} + \rho \frac{\nabla\mathcal{L}(w_{t})}{\|\nabla\mathcal{L}(w_{t})\|}$  not too far away from  $w_{t}$ . Hence, if GD does not diverge for this  $\eta$ , SAM does not either (unless the  $\rho$ -perturbation is non-negligible, i.e.,  $\| w_{t}\| \gg \rho$  no longer holds). This is not true for USAM: since the ascent step is un-normalized, it leads to a point far away from  $w_{t}$ , making the size of USAM updates much larger. In other words, the removal of normalization leads to much more aggressive steps, resulting in a different behavior than GD and also an easier divergence.
2.2 Generalizing to Non-Convex Cases: Scalar Factorization Problem
Now let us move on to non-convex losses. We consider a scalar version of the matrix factorization problem  $\min_{U,V}\frac{1}{2}\| UV^T -A\| _2^2$ , whose loss function is defined as  $\mathcal{L}(x,y) = \frac{1}{2} (xy)^2$ . Denote the initialization by  $(x_0,y_0)$ , then  $\mathcal{L}(x,y)$  is  $\beta \triangleq (x_0^2 +y_0^2)$ -smooth inside the region  $\{(x,y):x^2 +y^2\leq \beta \}$ . Hence, a learning rate  $\eta < 2 / \beta$  again allows GD to converge due to the well-known descent lemma. The following result compares the behavior of SAM and USAM under this setup.
Theorem 2 (Scalar Factorization Problem; Informal). For the loss function  $\mathcal{L}(x,y) = \frac{1}{2} (xy)^2$  restricted to a  $\beta$ -smooth region, if we set  $\eta = 1 / \beta < 2 / \beta$  (so  $GD$  finds a minimum), then
1. SAM. SAM never diverges and approaches a minimum within an  $O(\rho)$ -neighborhood (in fact, SAM with distinct  $\rho$ 's always find the same minimum  $(0,0)$ ).
2. USAM. On the other hand, USAM diverges once  $\rho \geq 15\eta$  - which could be much smaller than 1.
Thus, our observation in Theorem 1 is not limited to convex losses – for our non-convex scalar-factorization problem, the stability of SAM remains robust to the choice of  $\rho$ , while USAM is provably unstable. One may refer to Appendix D for the formal statement and proof of Theorem 2.
2.3 Experiment: Early-Stage Behaviors when Training Neural Networks
Figure 3: Behaviors of different algorithms when training neural networks  $(\eta = 0.025)$ .
As advertised, our result holds not only for convex or toy loss functions but also for practical neural networks. In Figure 3, we plot the early-stage behavior of GD, SAM, and USAM with different  $\rho$  values (while fixing  $\eta$ ). We pick two neural networks: a convolutional neural network (CNN) with tanh activation and a fully-connected network (FCN) with ReLU activation. We train them over the CIFAR10 dataset and report the early-stage training losses. Similar to Figure 1, Theorem 1 and Theorem 2, the stability of SAM is not sensitive to the choice of  $\rho$ , while USAM diverges easily.
3 Role of Normalization for Drifting Near Minima
Now, we explain the second role of normalization: enabling the algorithm to drift near minima. To convince why this is beneficial, we adopt a loss function recently considered by Ahn et al. (2023a) when understanding the behavior of GD with large learning rates. Their result suggests that GD needs a "large enough"  $\eta$  for enhanced performance, but this threshold can never be known a-priori in practice. To verify our observations from Figure 2, we study the dynamics of SAM and USAM over the same loss function and find that: i) no careful tuning is needed for SAM; instead, SAM with any configuration finds the same minimum (which is the "best" one according to Ahn et al. (2023a)); and ii) such property is only enjoyed by SAM – for USAM, careful tuning remains essential.
3.1 Toy Model: Single-Neuron Linear Network Model
To theoretically study the role of normalization near minima, we consider the simple two-dimensional non-convex loss  $\mathcal{L}(x,y)$  defined over all  $(x,y)\in \mathbb{R}^2$  as
$$
\mathcal {L} \colon (x, y) \mapsto \ell (x \times y), \quad \text {w h e r e} \ell \text {i s c o n v e x , e v e n , a n d 1 - L i p s c h i t z .} \tag {5}
$$
This  $\mathcal{L}$  was recently studied in (Ahn et al., 2023a) to understand the behavior of GD with large  $\eta$ 's. By direct calculation, the gradient and Hessian of  $\mathcal{L}$  at a given  $(x,y)$  can be written as:
$$
\nabla \mathcal {L} (x, y) = \ell^ {\prime} (x y) \left[ \begin{array}{l} y \\ x \end{array} \right], \quad \nabla^ {2} \mathcal {L} (x, y) = \ell^ {\prime \prime} (x y) \left[ \begin{array}{l} y \\ x \end{array} \right] ^ {\otimes 2} + \ell^ {\prime} (x y) \left[ \begin{array}{l l} 0 & 1 \\ 1 & 0 \end{array} \right]. \tag {6}
$$
Without loss of generality, one may assume  $\ell$  is minimized at 0 (see Appendix E for more details regarding  $\ell$ ). Then,  $\mathcal{L}$  achieves minimum at the entire  $x-$  and  $y$ -axes, making it a good toy model for studying the behavior of algorithms near a continuum of minima. Finally, note that the parametrization  $x\times y$  can be interpreted as a single-neuron linear network model - hence its name.
Before moving on to SAM and USAM we first briefly introduce the behavior of GD on such loss functions characterized in (Ahn et al., 2023a). Since  $\ell$  is even, without loss of generality, we always assume that the initialization  $w_{0} = (x_{0},y_{0})$  satisfies  $y_0\geq x_0 > 0$
Theorem 3 (Theorems 5 and 6 of (Ahn et al., 2023a); Informal). For any  $\eta = \gamma / (y_0^2 - x_0^2)$ , the  $GD$  trajectory over the loss function  $\mathcal{L}(x, y) = \ell(xy)$  has two possible limiting points:
1. If  $\gamma < 2$ , then the iterates converge to  $(0, y_{\infty})$  where  $y_{\infty}^{2} \in [\gamma/\eta - \mathcal{O}(\gamma) - \mathcal{O}(\eta/\gamma), \gamma/\eta + \mathcal{O}(\eta/\gamma)]$ .
2. If  $\gamma > 2$ , then the iterates converge to  $(0, y_{\infty})$  where  $y_{\infty}^{2} \in [2/\eta - \mathcal{O}(\eta), 2/\eta]$ .
Intuitively, the limiting point of GD (denoted by  $(0,y_{\infty})$  ) satisfies  $y_{\infty}^{2}\approx \min \{y_{0}^{2} - x_{0}^{2},2 / \eta \}$  . For simplicity, we denote  $\eta_{\mathrm{GD}}\approx 2 / y_0^2 -x_0^2$  as the threshold of  $\eta$  that distinguishes these two cases.
Interpretation of Ahn et al. (2023a). Fixing the initialization  $(x_0, y_0)$ , it turns out this model has a nice connection to the sparse coding problem, wherein it's desirable to get a smaller  $y_{\infty}^{2}$  (which we will briefly discuss in Section 4). According to Theorem 3, to get a smaller  $y_{\infty}^{2}$ , one must increase the learning rate  $\eta$  beyond  $\eta_{\mathrm{GD}}$ . Hence we mainly focus on the case where  $\eta > \eta_{\mathrm{GD}}$  - in which case we abbreviate  $y_{\infty}^{2} \approx 2 / \eta$  (see Table 1). However, GD diverges once  $\eta$  is too large - in their language,  $\gamma$  cannot be much larger than 2. This dilemma of tuning  $\eta$ , as we shall illustrate in Section 4 in more detail, makes GD a brittle choice for obtaining a better  $y_{\infty}^{2}$ .
On the other hand, from the numerical illustrations in Figure 4, one can see that SAM keeps moving along the manifold of minimizers (i.e., the  $y$ -axis) until the origin. This phenomenon is characterized in Theorem 4; in short, any moderate choice of  $\eta$  and  $\rho$  suffices to drive SAM toward the origin - no difficult tuning needed anymore!
In contrast, USAM does not keep moving along the axis. Instead, a lower bound on  $y_{\infty}^{2}$  also presents – although smaller than the GD version. As we will justify in Theorem 5, USAM does get trapped at some non-zero  $y_{\infty}^{2}$ . Thus, a dilemma similar to that of GD shows up: for enhanced performance, an aggressive  $(\eta, \rho)$  is needed; however, as we saw from Section 2, this easily results in a divergence.
Figure 4: Trajectories of different algorithms for the  $\ell(xy)$  loss ( $\eta = 0.4$  and  $\rho = 0.1$ ; initialization  $(x_0, y_0) = (2, \sqrt{40})$  is marked by a black dot).