Sharpness-Aware Minimization Efficiently Selects Flatter Minima Late In Training

Q9: Concerns about Assumptions 4.1-4.3.“I am not totally convinced that Assumptions 4.1-4.3 will be actually satisfied in practical situations. To be fair, many supporting references are provided in this paper.”

A9: Thank you for this fair comment. As noted, these assumptions are widely adopted in previous theoretical studies, and we have also cited many supporting references to demonstrate their relevance. Verifying whether these assumptions hold for deep neural networks is an interesting direction for future research and may have independent significance beyond our current focus.

Q10: Request for loss curves in more update steps. “I expected to see the loss going down in later steps in Figure 4 (b). Does the figure not show the entire steps?”

A10: Thank you for your question. Following your suggestion, we have extended the loss curve to include more update steps. In Figure R.4, we observe that the training loss initially increases exponentially fast and then gradually decreases, consistent with our theoretical predictions.

Q11: Overstated claims. “The result of Theorem 4.1 seems to be a little overstated … Theorem 4.1 does not show SAM converges although Theorem 4.3 does … Could the authors revise the sentence to a more accurate one?”

A11: Thank you for pointing out this problem. We have revised our claim P3 to “SAM selects a flatter minimum compared to SGD.”

Q12: “The paper uses the symbol $\approx$, but its mathematical definition is missing. Could the authors add the precise meaning?”

A12: Thank you for this question. First, we would like to emphasize that the symbol $\approx$ is not used in any formal theorems or proofs. In the main text, the symbol $\approx$ is used for two primary reasons:

• Ease of Understanding: The symbol $\approx$ simplifies the reading flow. For example, $\theta_0\approx\theta^\star$ in line 402 indicates that $\theta_0$ is close to $\theta^\star$, which is formally defined in Proposition 4.1.
• Approximate Observations: The symbol $\approx$ is also used to convey results that are not yet strictly established in the literature. For example, in line 362, the $\approx$​ is used to denote that the experimental phenomenon observed is approximate, as discussed in Wu et al. (2022), and has not yet been rigorously proven.

Following your suggestion, we have added additional commentary in the revised paper to explicitly explain the context in which $\approx$ is used.

Q13: Suggestion on the proof of Proposition 4.1. “Proof of Proposition 4.1 could be improved by adding comments on which part the authors used the assumptions.”

A13: Thank you for your suggestion. We have revised the proof of Proposition 4.1 to explicitly highlight where each assumption is applied. We believe this improvement will help readers follow the reasoning more easily.

Q14: “Does this refer to Eq. (3)? If so, it should be clearly stated: "Our experiments are conducted using the original implementation of SAM"”

A14: We appreciate your question. However, we are uncertain about the specific context that “this” refers to in your comment. Could you kindly provide further clarification or specify the exact lines where this concern arises? This will help us address your question accurately.

Reference

[1] Implicit Regularization in Matrix Factorization. In NeurIPS, 2017.

[2] Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In NeurIPS, 2018.

[3] Gradient Descent Provably Optimizes Over-parameterized Neural Networks. In ICLR, 2019.

[4] Computational-Statistical Gaps in Gaussian Single-Index Models. In COLT, 2024.

[5] Gradient Descent on Two-layer Nets: Margin Maximization and Simplicity Bias. In NeurIPS, 2021.

Q4: Technical problems. “how to obtain the inequality in lines 1137-1138; the phrase "which implies that $\boldsymbol{\theta}\_{T+1}=0$" in lines 1143-1144; the equality in lines 1085-1086; the inequality in lines 1087-1088.”

A4: Thank you for your careful review. We have addressed these concerns and correct the mistakes in the revised paper.

• Lines 1137-1138: The inequality can be derived by separately analyzing $\int\_{\theta\_{T+1/2}}^{2\theta\_T}\mathcal{L}''(z)\mathrm{d}z$ in two cases ($>0$ or $<0$) and applying two upper bounds of $\rho$ in each case. The detailed derivation has been included in our revised paper for clarity. Please kindly refer to our revised paper.
• Lines 1143-1144: We thank the reviewer for pointing this out. The phrase "which implies that $\boldsymbol{\theta}\_{T+1}=0$" has been deleted in the revised version.
• Lines 1085-1086: This equality follows from $\boldsymbol{\theta}+\rho\nabla\mathcal{L}(\boldsymbol{\theta})=(I+\rho G) \boldsymbol{\theta}$, given $\mathcal{L}(\boldsymbol{\theta})=\frac{1}{2}\boldsymbol{\theta}^\top G \boldsymbol{\theta}$).
• Lines 1087-1088: The inequality holds because $G=\frac{1}{n}\sum\_{i=1}^n\nabla f(\boldsymbol{x}_i;\boldsymbol{\theta})\nabla f(\boldsymbol{x}_i;\boldsymbol{\theta})^\top\succeq0$. Thus, $(I+\rho G)G(I+\rho G)=G+2\rho G^2+\rho^2 G^3 \succeq G$.

Q5: “Theorems 4.1 and 4.2 rely on the linear approximation (5), so it does not explain the whole dynamics after escaping from a sharp minimum and before getting close to a flat minimum. Furthermore, combined with the squared loss, the analysis is in a classical strongly convex setup.”

A5: Thank you for your thoughtful comments.

• Beyond the dynamics near the minima: Theorems 4.1 and 4.2 are indeed based on linear stability and apply to the dynamics near the minima. However, it is worth noting that Proposition 4.1 and Theorem 4.3 are not based on linear approximation. In particular, Proposition 4.1 addresses the high-loss region where the linear approximation does not hold (as detailed in the beginning of Section 4.2). Together, these four results provide a more comprehensive view of the training dynamics.
• Inquiry for “whole” training dynamics: Characterizing the "whole" training dynamics is inherently challenging under general settings. Most existing analyses of the "whole" dynamics, even for (S)GD, are limited to extremely simplified setups, such as single-index models [4] or feature learning regime [5].
• Convex setup, not strongly convex. Finally, we clarify that our analysis in Theorems 4.1 and 4.2 assumes a convex setup rather than a strongly convex one. This distinction is crucial because over-parameterized models typically exhibit many degenerate directions, which we discuss in more detail in Lemma C.2.

Q6: Concerns about Example 4.1. “The single-data setup of Example 4.1 is unrealistic, and I do not know how relevant it is.”

A6: We would like to re-clarify that Example 4.1 is primarily intended to clearly illustrate our two-phase picture. While this example is simple, it is representative in several key ways: 1) it uses a two-layer neural network; 2) addresses a fitting problem; 3) includes a landscape with multiple valleys; 4) satisfies the sub-quadratic property; 5) exhibits differences in flatness between minima; 6) and highlights SAM’s two-phase training behavior.

In Sec. 4.1-4.3, we present our formal theoretical results in general case with corresponding experimental evidence to support our two-phase picture.

Q7: Concerns about Proposition 4.1. “Proposition 4.2 considers the case in which there is only one parameter and has almost no practical usefulness.”

A7: Thank you for this valuable comment. In fact, generalizing Proposition 4.1 beyond $p=1$ presents significant challenges. Extending Proposition 4.1 to higher dimensions requires a more comprehensive definition of sub-quadratic properties in high-dimensional settings. This necessitates a deeper exploration of the loss landscape, which lies beyond the scope of our primary focus on SAM's dynamics and is left for future work.

Q8: “Theorem 4.2 directly implies Theorem 4.1, so the latter should be a corollary.”

A8: Thank you for pointing out this problem. In fact, the proof of Theorem 4.2 builds upon the results of Theorem 4.1. To reflect this dependency, we have revised Theorem 4.2 into Corollary 4.2.

Thank you for your great efforts on the review of this paper and for appreciating our recognizing the value of our theoretical analyses. We will try our best to address your questions.

If you still have further concerns, or if you are not satisfied by the current responses, please let us know, so that we can update the response ASAP.

Q1: “Eq. (4) and the update rule that the authors use in theory (the two-step update rule defined at the beginning of C.1) are not equivalent.”

A1: Thank you for your careful review. We acknowledge that there is indeed a difference between these two formulations. Specifically, the update rule used in Appendix C.1 assumes independence between the randomness in the inner and outer updates of SAM. This simplification is a common technical approach, as decoupling the randomness in SAM's updates is challenging. Notably, this technique has also been employed in the classic work on understanding SAM (see Appendix C.2 of Andriushchenko & Flammarion (2022)). We agree that analyzing the original formulation of SAM is an intriguing direction for future research.

Q2: Concerns about our theoretical setup regarding the update rule. “The experiments and the theory part use different methods (Eqs. (3) and (4)).” and “Are there any experiments using Eq. (4)?”

A2: We are glad to discuss more on the discrepancy between our experimental and theoretical settings. First, we’d like to emphasize that the simplification of approximating the step size $\rho_t$ with a constant $\rho$ has been empirically justified by Andriushchenko & Flammarion (2022), and has been widely adopted in recent theoretical analysis of SAM (see more references in Sec. 2). This simplification maintains mathematical tractability without significantly affecting the general behavior of SAM, which has been supported by empirical evidence.

Additional Experiments: Second, we have conducted new experiments using Eq. (4). Here, we refer to the update rule in Eq. (4) as Unnormalized-SAM (USAM). Specifically, we measured how the generalization gap between $\boldsymbol{\theta}\_{{\rm SGD \to USAM}, t}^T$ and $\boldsymbol{\theta}\_{{\rm USAM}}^T$, as well as the sharpness of $\boldsymbol{\theta}\_{{\rm SGD \to USAM}, t}^T$, changes as the USAM training proportion $\frac{T-t}{T}$ increases. In Figure R.5 (please click the link to check the figure), both the generalization gap and sharpness drop sharply once the USAM training proportion exceeds zero, consistent with our main discovery in Fig. 3. This finding validates the effectiveness of USAM in capturing the main behavior of SAM, further bridging the gap between our theoretical results and empirical findings.

Q3: Concerns about our theoretical setup regarding the loss function. “This paper only studies the squared loss.”

A3: Thank you for your comment. The use of the squared loss function is a common practice in theoretical analyses of neural networks, and it facilitates the study of various aspects including implicit bias [1], training dynamics [2], and convergence [3]. While this choice simplifies our theoretical framework, our theoretical results remain representative for understanding the training dynamics and generalization behavior of neural networks.

Additional Experiments: To further address your concerns, we have conducted additional experiments to reproduce our main discovery in Sec. 3 using the squared loss. Consistently with the main paper, we measured how the generalization gap between $\boldsymbol{\theta}\_{{\rm SGD \to SAM}, t}^T$ and $\boldsymbol{\theta}\_{{\rm SAM}}^T$, as well as the sharpness of $\boldsymbol{\theta}\_{{\rm SGD \to SAM}, t}^T$, changes as the SAM training proportion $\frac{T-t}{T}$ increases. In Figure R.6, both the generalization gap and sharpness follow trends consistent with our previous findings in Fig. 3, thereby bridging the gap between our theoretical and experimental results.

Thank you for the quick response! Yes, it refers to Eq. (3). We will revise line 138 into "Our experiments are conducted using the original implementation of SAM (see Eq. (3))."

Yes, we clarify that $\mathbb{E}$ is taken with respect to the random variables $\xi_t^1$, $\xi_t^2$ ($t\geq0$), which are the indices of independently random examples sampled from $\mathcal{D}_{\mathrm{train}}$.

We apologize for the potential confusion caused by the notation and re-clarify it: $\boldsymbol{\xi}\^B(\boldsymbol{\theta}\^{t+1/2}):=\mathcal{L}\_{\xi\_t\^2}(\boldsymbol{\theta}\^{t+1/2})-\nabla\mathcal{L}(\boldsymbol{\theta}\^{t+1/2})$. Observing that $\boldsymbol{\theta}^t-\eta \nabla \mathcal{L}(\boldsymbol{\theta}^{t+1/2})=\phi(\xi\_1\^1,\xi\_1\^2,\cdots,\xi_t\^1)$ and $\boldsymbol{\theta}^{t+1/2}=\psi(\xi\_1\^1,\xi\_1\^2,\cdots,\xi_t\^1)$, we can use the law of total expectation to derive

PS: If no further concerns remain, we will really appreciate it if you could update the confidence score to reflect those resolutions in our responses. Thank you!

Q6: “What are the advantages of SAM shown in Theorem 4.3? SGD converges at the same rate under the Assumptions 4.2 and 4.3.”

A6: Thank you for this insightful question.

• Clarification on Theorem 4.3's purpose: First, we would like to clarify that Theorem 4.3 is not intended to demonstrate the advantages of SAM over SGD. Instead, the advantage of SAM is primarily established in Theorem 4.1, which shows that SAM selects flatter minima than SGD, leading to better generalization. Instead, Theorem 4.3 aims to support "P4: The convergence rate of SAM is extremely fast" in Phase II (see Tab. 1). Notably, Theorem 4.3 establishes a faster convergence rate of SAM compared to the classic result from Andriushchenko & Flammarion, (2022).

• Comparison to SGD: As noted by the reviewer, under the same conditions, SGD and SAM exhibit the same convergence rate. Indeed, there is currently no consistent empirical evidence suggesting that SAM outperforms SGD in terms of convergence rate.

Thank you for your great efforts on the review of this paper and for your constructive comments on our work. We will try our best to address your questions.

If you still have further concerns, or if you are not satisfied by the current responses, please let us know, so that we can update the response ASAP.

Q1: Suggestions for clarifications of Definition 4.1 (Linear Stability). “Does it mean that SAM is used to optimize linearized function but is evaluated on the real loss or that SAM both trains and tests on the linearized loss?” and “Can the authors give a clear definition of linear stability?”

A1: Thank you for this valuable suggestion.

• Clarification on the use of linearized loss: First, we’d like to clarify that the linear stability theory, which involves the linearized loss, is used solely for theoretical analysis. All of our experiments are conducted on the actual training and testing loss, not the linearized version.

• Definition of linear stability: Second, linear stability theory studies the training dynamics of SAM near the global minima of training loss. Specifically, near the global minimum $\boldsymbol{\theta}^\star$, the model $f$ can be approximated by its linearization $f\_{\rm lin}$ (see Eq. (5)). Consequently, the training loss can be expressed as $\mathcal{L}(\boldsymbol{\theta})=\frac{1}{2n}\sum\_{i=1}^n(f\_{\rm lin}(\boldsymbol{x}_i;\boldsymbol{\theta})-y_i)^2=\frac{1}{2}(\boldsymbol{\theta}-\boldsymbol{\theta}^\star)^\top \Big(\frac{1}{n}\sum\_{i=1}^n \nabla f(\boldsymbol{x}_i;\boldsymbol{\theta}^\star)\nabla f(\boldsymbol{x}_i;\boldsymbol{\theta}^\star)^\top\Big) (\boldsymbol{\theta}-\boldsymbol{\theta}^\star)$. The linear stability analysis focuses on SAM's dynamics (Eq. (4)) when optimizing this linearized loss $\mathcal{L}(\boldsymbol{\theta})$​. This theoretical analysis demonstrates the flatness of the minima selected by SAM, which correlates with its generalization capability.

In addition, following your suggestion, we have further clarified Definition 4.1 in our revised paper.

Q2: “What is $a$ in Proposition 4.1, which is used to give an upper bound of $\rho$?”

A2: Thanks for this inquiry. $a$ is defined in Definition 4.2 as the sharpness $\mathcal{L}''(\theta^\star)$ of the global minimum $\theta^\star$.

Q3: Concerns on Theorem 4.2. “Is the first claim in Theorem 4.2 the contrapositive of Theorem 4.1? If not, can you explain the difference between these two claims?”

A3: Thank you for this insightful question. Yes, Theorem 4.2 is indeed the contrapositive of Theorem 4.1. To clarify this relationship, we have revised Theorem 4.2 and rephrased it as Corollary 4.2 in the revised paper.

Q4: “After Theorem 4.1 and Theorem 4.2, the authors claim that SAM escapes from the sharp minima exponentially fast, which sounds like a good result. However, in Theorem 4.2, the result says that the final loss is $C^t$ larger than the initial one, which seems to be a bad result. What is the relationship of these two results or the exponentially fast claim should be derived from some other theorem?”

A4: Thank you for raising this concern.

• Definition of Exponentially Fast Escape: We clarify that the definition of exponentially fast escape is precisely expressed as $\mathbb{E}[\mathcal{L}(\boldsymbol{\theta}_t)]\geq C^t \mathbb{E}[\mathcal{L}(\boldsymbol{\theta}_0)]$, where $C>1$. This reflects the rapid increase in loss during the escape phase, as captured in Theorem 4.2 (now Corollary 4.2).
• Escape then converge: The exponential escape in Theorem 4.2 applies only to Phase I. Once SAM transitions into Phase II, the loss begins to descend, and SAM converges to a flatter minimum.

Q5: Concerns about Proposition 4.1. “When the authors want to prove that SAM can stay in the same valley as SGD, there is only one valley in Proposition 4.1. Thus, to show SAM stays in the same valley, the function should have at least two valleys.” and “Can the authors provide proof when the function has two valleys, SAM will not go from one valley to another?”

A5: Thank you for this insightful question. We’d like to clarify that multiple valleys can indeed exist in the setting of Proposition 4.1:

• Proposition 4.1 assumes the landscape is sub-quadratic within the valley $V=[-2b,2b]$, but it does not impose any assumptions about the landscape outside $V$. Therefore, additional valleys can exist in $(-\infty,-2b)\cup(2b,+\infty)$.
• For example, consider $\mathcal{L}(\theta)=1-\cos(\frac{\pi\theta}{2b})$, which satisfies the sub-quadratic property within $V=[-2b,2b]$. This loss function has infinitely many valleys: $[2(2k-1)b,2(2k+1)b],k\in\mathbb{Z}$. By Proposition 4.1, SAM remains in the current valley $V=[-2b,2b]$ and cannot enter into other valleys.

We have added this clarification to the revised paper to address potential misunderstandings.

Dear Reviewer cw28,

We hope that our responses could adequately address your concerns. As the deadline of this discussion phase is approaching, we warmly welcome further discussion regarding any additional concerns that you may have, and we sincerely hope you can reconsider the rating accordingly.

Thank you for the time and appreciation that you have dedicated to our work.

Best regards,

Authors of submission 1483

Thank you for your great efforts on the review of this paper and for appreciating our novelty and contributions! We will try our best to address your questions.

If you still have further concerns, or if you are not satisfied by the current responses, please let us know, so that we can update the response ASAP.

Q1: “For some experiments (Fig 9) using SGD toward the end of training outperformed SAM in terms of accuracy, which is mildly contradicting toward an implicit claim that at the end-of-training, SAM outperforms SGD. Some commentary in the main section on these results would be appreciated.”

A1: Thank you for your insightful suggestion. In our revised paper, we have included additional commentary in the main section to address the mild contradiction you mentioned. We acknowledge that in Fig. 9, when SAM training proportion approaches one, switching from SAM to SGD may slightly outperform full SAM training in terms of accuracy.

• However, we would like to point out that accuracy is a discrete variable and can exhibit small fluctuations, making it a less stable measure. For example, in Fig. 9 for ResNet-20 on CIFAR-10, the switching method only outperforms full SAM training by approximately $10^{-4}$ in accuracy, which we consider to be a minor fluctuation.
• In contrast, when examining test loss, we observe consistent behavior, with full SAM training continuing to outperform using SGD toward towards the end of training.

Q2: Concerns about Proposition 4.1. “The size of the valley in Proposition 4.1 for P2, (−2b,2b), appear extremely wide when the parameter is initialized between (−b,b), which weakens the significance of this proof.”

A2: Thank you for this constructive comment. In fact, the result can be extended to almost all initializations within the valley $(-2b,2b)$. Specifically, we can prove that a similar result to Proposition 4.1 holds for all initializations within $(-(2-\epsilon)b,(2-\epsilon)b), \forall \epsilon\in(0,1)$. By choosing $\epsilon$ sufficiently small, the result effectively applies to nearly all initializations within $(-2b,2b)$. We have included this generalized result for Proposition 4.1 in the revised paper.

Q3: Suggestions for more general loss landscape visualization. “The empirical evidence (Fig 4c) for this claim also appears limited. Can a more general visualization of the loss landscape be considered?”

A3: Thank you for your suggestion. We’d like to clarify that the one-dimensional linear interpolation presented in Fig. 4 (c) is widely adopted for visualizing the complex loss landscape of neural networks [1-3], and it is sufficient to support our key claim P2 that the SAM remains within the current valley during escape.

Additional Experiments: To provide a more general visualization, we have conducted additional experiments using a two-dimensional visualization approach, as proposed by [4]. In this setup, we center the visualization at $\boldsymbol{\theta}\_{{\rm SGD \to SAM}, t}^T$ and select two directions: $\boldsymbol{\delta}\_{\rm target}$ and $\boldsymbol{\delta}\_{\rm random}$. Specifically, $\boldsymbol{\delta}\_{\rm target} = \boldsymbol{\theta}\_{{\rm SGD}}^T - \boldsymbol{\theta}\_{{\rm SGD \to SAM}, t}^T$, pointing towards the SGD solution, and $\boldsymbol{\delta}\_{\rm random}$ is a random direction sampled from $\mathcal{N}(0, \boldsymbol{I}_p)$ , with layer-wise normalization [4] applied. We then plot the loss landscape as a function $f(\alpha, \beta) = \mathcal{L}(\boldsymbol{\theta}\_{{\rm SGD \to SAM}, t}^T + \alpha \boldsymbol{\delta}\_{\rm target} + \beta \boldsymbol{\delta}\_{\rm random})$. In Figure R.3 (please click this link to check the figure), it is clear that $\boldsymbol{\theta}\_{{\rm SGD \to SAM}, t}^T$ and $\boldsymbol{\theta}\_{{\rm SGD}}^T$ stay within the same valley. Moreover, the contours around $\boldsymbol{\theta}\_{{\rm SGD \to SAM}, t}^T$ is wider compared to $\boldsymbol{\theta}\_{{\rm SGD}}^T$, indicating $\boldsymbol{\theta}\_{{\rm SGD \to SAM}, t}^T$ corresponds to a flatter minimum. This two-dimensional approach provides additional insights into the structure of the loss landscape and further supports our two key claims P2 and P3.

Q4: “In 'Support for P3', can you support why increasing the loss might lead to flatter minima and better generalization? The empirical evidence suggests that SAM decreases the training loss, in addition to test loss.”

A4: We suspect you are asking about "the ascent update $\rho\nabla\mathcal{L}(\boldsymbol{\theta}_t)$ tends to increase the loss" in L348-349 in our "Support for P3".

• Intuitive Explanation: Intuitively, the inner update $\boldsymbol{\theta}\_{t+1/2}=\boldsymbol{\theta}_t+\rho\nabla\mathcal{L}(\boldsymbol{\theta}_t)$ tends to increase the loss, making SAM more unstable (and more prone to divergence) than SGD. As a result, SAM must identify flatter minima to satisfy its stricter stability condition. Please kindly refer to Appendix C.3 for the formal proof.
• Inner vs. Outer Updates: Notice that this claim about loss increase only pertains to the inner update of SAM, i.e., the comparison between $\mathcal{L}(\boldsymbol{\theta}_t)$ and $\mathcal{L}(\boldsymbol{\theta}\_{t+1/2})$. This is not inconsistent with observed outer loss decrease, i.e., the comparison between $\mathcal{L}(\boldsymbol{\theta}\_{t})$ and $\mathcal{L}(\boldsymbol{\theta}\_{t+1})$. In fact, our theory (see Theorem 4.3) and experiments demonstrate the convergence and loss descent of SAM.

Reference

[1] Qualitatively Characterizing Neural Network Optimization Problems. In ICLR, 2015.

[2] On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima. In ICLR, 2017.

[3] Linear Mode Connectivity and the Lottery Ticket Hypothesis. In ICML, 2020.

[4] Visualizing the Loss Landscape of Neural Nets. In NeurIPS, 2018.

Thank you for your great efforts on the review of this paper and for appreciating our novelty and contributions! We will try our best to address your questions.

If you still have further concerns, or if you are not satisfied by the current responses, please let us know, so that we can update the response ASAP.

Q1: Concerns about the limitation of our theory regarding the choice of loss function. “The divergence between paper's experimental and theoretical setting is notable especially in the choice of the loss function. This is acknowledged in the paper and is acceptable given the novelty of their results.”

A1: Thank you for appreciating the novelty of our work. The use of the squared loss function is a common practice in theoretical analyses of neural networks, and it facilitates the study of various aspects including implicit bias [1], training dynamics [2], and convergence [3]. While this choice simplifies our theoretical framework, our theoretical results remain representative for understanding the training dynamics and generalization behavior of neural networks.

Additional Experiments: To further close the gap between our experimental and theoretical settings, we have conducted additional experiments to reproduce our main discovery in Sec. 3 using the squared loss. Consistently with the main paper, we measured how the generalization gap between $\boldsymbol{\theta}\_{{\rm SGD \to SAM}, t}^T$ and $\boldsymbol{\theta}\_{{\rm SAM}}^T$, as well as the sharpness of $\boldsymbol{\theta}\_{{\rm SGD \to SAM}, t}^T$, changes as the SAM training proportion $\frac{T-t}{T}$ increases. In Figure R.6 (please click the link to check the figure), both the generalization gap and sharpness follow trends consistent with our previous findings in Fig. 3, thereby bridging the gap between our theoretical and experimental results.

Q2: Writing problems. “Although the paper is written well overall, certain unclearly defined concepts, choice of notation, and design choices make it hard to follow at times. I provide more detailed feedback regarding these in the section below.”

A2: Thank you for your careful review. Your feedback is invaluable in helping us improve our paper. We have carefully revised the paper to address the writing problems you highlighted. Below, we outline the key revisions made:

• Clarification of Concepts (L000, L115): Additional explanations and definitions to ensure clarity and accessibility.
• Improved Notation (L110, L286, L352): Streamlined notations for better consistency and readability.
• Corrections (L177, L291, L354): Fixed typos and other minor issues.

Regarding the other specific points you raised:

• L081: We kindly request further elaboration on this issue to ensure we fully understand your concern and address it appropriately.
• L177 (X-axes): To maintain consistency across Fig. 3, 5 and 6-9 and facilitate comparison, we have decided to retain the current x-axis formatting.
• L318: Following the notation $\ell\_i(\boldsymbol{\theta})$, which denotes the loss on the $i$-th sample (see Line 113), $\nabla \ell\_{\xi}(\boldsymbol{\theta})$ represents the 1-batch loss gradient (at $\boldsymbol{\theta}$), where $\xi$ denotes the index of a random sample.

We hope our revision and explanations can address your concerns. Thank you again for your constructive feedback.

Q3: Questions about Proposition 4.1. “L388: Do these results generalize to cases where $p>1$ and minibatch SAM?”

A3: Thank you for the insightful question. In fact, generalizing Proposition 4.1 beyond $p=1$ and full-batch SAM presents significant challenges:

• Extending to $p>1$: Generalizing Proposition 4.1 to higher dimensions would require a more comprehensive definition of sub-quadratic properties in high-dimensional settings. This necessitates a deeper exploration of the loss landscape, which lies beyond the scope of our primary focus on SAM's dynamics and is left for future work.
• Mini-batch SAM: The sub-quadratic property pertains to the overall loss $\mathcal{L}(\boldsymbol{\theta})$ across the entire training set and does not impose assumptions on the mini-batch loss $\ell_i(\boldsymbol{\theta})$. Therefore, extending our analysis to mini-batch SAM would require characterizing the landscape of $\ell_i(\boldsymbol{\theta})$, which introduces additional complexities and is not directly addressed in our current framework.

Q4: Questions about the sub-quadratic property. “L406: Can we assume any valley around $\boldsymbol{\theta}^\*$ found by SGD to be sub-quadratic?”

A4: Thank you for your question. As mentioned above, research on the sub-quadratic property is still in its early stages, and its definition in high-dimensional settings remains unclear. Answering this question will require a deeper investigation into the structure of the loss landscape.

Q5: Questions about the robustness-accuracy trade-off. “L524: Are you making a claim about robustness-accuracy trade-off? I.e. do you expect robustness to increase without sacrificing vanilla accuracy if the AT is started late?”

A5: Thank you for your question.

• We are not making a claim about the robustness-accuracy trade-off in this paper, as it is not the primary focus of our work. Our experiments in Sec. 5 aim to show that Adversarial Training (AT) can efficiently improve robustness even when applied only in the final few epochs.
• Regarding the second question, we do observe a slight drop in accuracy on clean samples when starting AT late in training. This trade-off is expected and aligns with findings reported in the literature [4, 5].

Reference

[1] Implicit Regularization in Matrix Factorization. In NeurIPS, 2017.

[2] Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In NeurIPS, 2018.

[3] Gradient Descent Provably Optimizes Over-parameterized Neural Networks. In ICLR, 2019.

[4] Robustness May Be at Odds With Accuracy. In ICLR, 2019.

[5] Theoretically Principled Trade-off between Robustness and Accuracy. In ICML, 2019

Thank you for your great efforts on the review of this paper and for recognizing the value of our contributions. We will try our best to address your questions.

If you still have further concerns, or if you are not satisfied by the current responses, please let us know, so that we can update the response ASAP.

Q1: Concerns about practical applicability of our theoretical analysis. “They do not sufficiently explain why transitioning the training method from SGD to SAM results in a large generalization gap compared to using the full SAM training method. The experiments in the paper demonstrate that switching from SAM to SGD often leads to poorer results.”

A1: Thank you for your question. In Fig. 3, switching from SGD to SAM does not result in a large generalization gap compared to full SAM training. Based on your follow-up comment, we believe your concern may be about the practical applicability of our theoretical analysis to the case of switching from SAM to SGD.

• Clarification of the purpose of our theoretical analyses: First, we would like to clarify that the primary goal of our theoretical analysis is to explain why SAM efficiently finds flatter minima compared to SGD when applied in the late stage of training. Our theoretical setup closely reflects the real-world settings, and in Fig. 4, our theoretical claims are well-supported by experiments conducted in practical scenarios. Therefore, our theoretical results are indeed applicable to real training scenarios and sufficiently explain our main discovery.

• Insights into the SAM-to-SGD: Second, while our theoretical results do not directly explain the empirical observations when switching from SAM to SGD, they offer valuable insights into the underlying training behavior during this transition.

  • In our theory (see Theorem 4.1 and Tab. 2), the global minima selected by SAM are inherently linearly stable for SGD. However, linear stability is a necessary but not sufficient condition for minima selection. It is possible that other conditions exist and drive SGD escapes the flatter minima found by SAM, resulting in relatively sharper minima and thus poorer generalization. Nevertheless, the relatively sharper minima are still required to be linear stable for SGD.
  • In Sec. B.4.1, we explore how sharpness at the iterator evolves during training when we switches from SAM to GD. In Fig. 10, we observe that the sharpness rapidly increases once the switch occurs and then oscillates around $2/\eta$ (the linear stability condition for GD (Wu et al. 2018), also commonly known as the Edge of Stability phenomenon), regardless of how late the switch occurs. This implies that SGD immediately escapes the flat minima found by SAM and converges to relatively sharper minima, which aligns with the empirical observations in Sec. 5.

  Together with our theory, we conjecture that other than linear stability, additional conditions govern the minima selection process and drive (S)GD towards relatively sharper solutions. We will leave the exploration of the conditions beyond linear stability as future directions.

Q2: Misunderstandings regarding our main contribution. “In theory, the switch from SGD to SAM is expected to yield certain benefits; however, no experimental evidence is provided to support this.”

A2: We'd like to re-clarify the main contributions and logic progression of this paper:

1. Sec. 3: we empirically find that the solutions obtained by switching from SGD to SAM yield comparable generalization and sharpness than solutions found by full SAM training, even when the switch occurs a few epochs before the end of training.
2. Sec. 4: We theoretically establish a two-phase picture to explain the empirical findings. This picture is characterized by four key claims, all of which are well-supported by practical experiments.

Together, the experimental evidence in Sec. 3 demonstrate the benefits of switching from SGD to SAM in terms of generalization and sharpness, and our theory in Sec. 4 provides a sufficient explanation for these benefits.

Q3: Suggestions for empirical evidence on transformer-based models. “The models discussed are primarily convolution-based, so it would be beneficial to include some transformer-based models as well.”

A3: Thank you for your suggestion. We have conducted new experiments on Vision Transformers (ViTs) to further verify that applying SAM in the last few epochs achieves generalization comparable to full SAM training. As Adam is a more common practice for training ViTs, we replace SGD with Adam before the switch. Then, consistently with the main paper, we measured how the generalization gap between $\boldsymbol{\theta}\_{{\rm Adam \to SAM}, t}^T$ and $\boldsymbol{\theta}\_{{\rm SAM}}^T$ changes as the SAM training proportion $\frac{T-t}{T}$ increases. As shown in Figure R.1 (please click the link to check the figure), the generalization gap approaches zero once the SAM training proportion exceeds zero. This further strengthens our argument, confirming the efficiency of SAM during the late phase of training.

Q4: Question about the effect of perturbation radius when switching from SAM to SGD. “Could the authors explain how the neighborhood size affects the results when switching from SAM to SGD?”

A4: An interesting question. In our theory (see Theorem 4.1 and Tab. 2), the perturbation radius $\rho$ influences the sharpness of the global minima selected by SAM. Specifically, a larger perturbation radius leads to the selection of flatter global minima. In contrast, SGD's behavior is independent of the perturbation radius. However, as discussed in our conjecture in Sec. 5, the properties of the final solutions are primarily shaped by the optimization method chosen in the late training phase. Consequently, when switching from SAM to SGD, the generalization properties of the final solution are expected to be independent of the perturbation radius, as SGD dominate in the final stage of training.

Additional Experiments: To further explore this, we have conducted new experiments to investigate the effect of perturbation radius on generalization when switching from SAM to SGD. Specifically, we varied the perturbation radius $\rho \in \\{0.02, 0.05, 0.1, 0.15, 0.2\\}$ for SAM and measured how the test loss and error of $\boldsymbol{\theta}\_{{\rm SAM \to SGD}, t}^T$ change as the SAM training proportion $\frac{t}{T}$ increases. In Figure R.2, we observe no significant differences in test loss or error across different $\rho$ values, even when training with SAM for most of the time (e.g., $\frac{t}{T} = 0.8$). Notably, only full SAM training (i.e., $\frac{t}{T}=1$) shows a pronounced dependence on $\rho$, where a larger perturbation radius leads to better generalization. These findings align with our expectation that the perturbation radius has minimal impact on generalization when transitioning from SAM to SGD, further validating our theoretical results in Sec. 4 and conjecture in Sec. 5.

Hi Zhiwei,

Thank you for your appreciation of the novelty and significance of our study!

We reviewed the paper you mentioned, which introduces a sharpness-regularized training algorithm with a formulation similar to SAM. Notably, its ablation study demonstrates the algorithm's effectiveness when applied in the mid or late training stages. We will include a citation to this paper in our revision.

Best regards,

Authors of submission 1483

Dear AC and reviewers,

We sincerely thank you for your valuable feedback and constructive suggestions. Following the recommendations from AC and Reviewer S3eC, we have further revised our paper. The additional changes are highlighted in blue. Below is a summary of the updates:

• Additional acknowledgement of Andriushchenko et al. (2022): As suggested by AC, we have included an additional commentary in the introduction section to explicitly acknowledge that Andriushchenko et al. (2022) made a similar finding.
• Clarifications in the proof of Theorem 4.1: As suggested by Reviewer S3eC, we have
  • Added Remark C.1 in Sec. C.1 to clarify the use of $\mathbb{E}$.
  • Added Remark C.2 to explain that $\boldsymbol{\xi}^B(\boldsymbol{\theta}^{t+1/2})$ and $\boldsymbol{\theta}^t-\eta \nabla \mathcal{L}(\boldsymbol{\theta}^{t+1/2})$ are uncorrelated.

Since authors cannot upload the revised paper directly to OpenReview, you may access the updated version here.

We hope these updates address the concerns raised.

Best regards,

Authors of submission 1483

Dear AC,

Thank you for your question. We indeed acknowledge in a separate paragraph of our paper (lines 221-224) that Andriushchenko et al. (2022) made a similar observation, but they did not give it significant attention. Our work isolates and extends this finding by showing that even fewer epochs of SAM (e.g., around 5 epochs for WideResNet-16-8 on CIFAR-10) applied at the end of training could achieve comparable test error and loss to full SAM training (see Fig. 3). More importantly, as noted by Reviewer 6K1k, our study provides an in-depth examination of SAM’s late phase behavior, both empirically and theoretically, an area that has been understudied in prior research. Here, we provide a concise summary of our extra contributions:

• More empirical investigations.

  • New experiments regarding the sharpness. Other than generalization ability, we investigate the effect of late-phase SAM on the sharpness of the final solution. This experiment directly verifies the implicit bias of SAM towards flatter minima when applied in the late phase of training.
  • Extensive experiments under various setups. We conduct extensive experiments under various settings to justify the efficiency of late-phase SAM in selecting flatter minima. Our experiments cover various variants of SAM (e.g. vanilla SAM, ASAM and USAM), loss function (e.g., squared loss and cross entropy loss), architectures (e.g. convolution-based and transformer-based models), and datasets.

• Rigorous theoretical justifications. Our theoretical analysis provably uncovers a two-phase picture in the training dynamics after switching to SAM in the late phase, which highlights the distinct value of our work, as acknowledge by Reviewer WFBm, bzn5 and S3eC.

  • This two-phase picture is characterized by four key claims in Tab. 1, each rigorously supported by corresponding theorems or propositions. For example, our Corollary 4.2 explicitly prove the escaping behavior of SAM after the switch, which is also suggested by Andriushchenko et al. (2022). However, they did not provide an explicit proof or direct empirical evidence for this claim.
  • We also design careful experiments in real-world setups to verify each claim (see Fig. 4, 14, and 16).

• In-depth discussions and extensions. Beyond our central findings, we also investigate the necessity of using SAM in the early phase. We find that early-phase SAM may offer only limited improvements over SGD in terms of generalization and sharpness (see Fig. 5 and 15). Particularly, in Fig. 10, we observe that the sharpness rapidly increases once switching from SAM to GD and then oscillates around $2/\eta$​ (commonly known as the Edge of Stability phenomenon).

  • Conjecture on late training phase. Together, contrary to the conventional view that early training is critical (Achille et al. 2019; Frankle et al. 2020), we conjecture that the optimization algorithm chosen at the end of training is more critical in shaping the final solutions’ properties.
  • Extension from SAM to AT. We also validate our conjecture by extending our finding on SAM to Adversarial Training (AT). Specifically, we observe that AT efficiently enhances the robustness of neural networks when applied late in training.

In summary, we believe our work provides distinct and significant contributions over Andriushchenko et al. (2022).

We hope this clarification can address your question.

Best,

Authors of submission 1483