A Unified Stability Analysis of SAM vs SGD: Role of Data Coherence and Emergence of Simplicity Bias

We thank the reviewer for the response. Here we provide brief summary for our paper to help understanding it in more depth and we are glad to answer any question that the reviewer may have: This paper is for the first time studies what properties of data (and not just global properties of the function) ensure stochastic optimization algorithms like Stochastic Gradient Descent (SGD) and Sharpness-Aware Minimization (SAM) settle in or escape from local minima in complex NN training, and why some local minima they find generalize better than others.

At the heart of our work is the idea of data coherence — a new way to measure how consistent the gradients (training signals) are across different data points. We show that when the gradients point in more similar directions (i.e., high coherence), optimization tends to be more stable, and the resulting models tend to generalize better. SAM is better at finding such minima, and SGD will `settle' in some incoherent minima that SAM will escape from. We analyze this behavior mathematically in two layer ReLU networks. Using this framework, we prove that stability during training is influenced by both the data and the algorithm, and we show that SAM encourages the algorithm to prefer more coherent, stable solutions than vanilla SGD.

For the novelty of our theory, we summarize as follows:

1. Role of coherence in the optimization: We highlight the role of data coherence in shaping optimization dynamics, showing for the first time the importance of distribution of data itself (beyond its global properties such as Lipschitz-ness) and hot it actively interacts with the geometry of the loss landscape. In particular, coherent data distributions lead to more stable optimization trajectories, while datasets with many outliers (i.e., low coherence) can exhibit instability. This perspective helps bridge the gap between optimization theory and data.

2. New analysis for SAM: Our framework provides a novel analysis of Sharpness-Aware Minimization (SAM), showing that SAM implicitly penalizes incoherent solutions. This reveals that SAM not only promotes flat minima but also enforces stronger alignment between the solution and the underlying data, offering new insight into its generalization behavior.

3. We show that coherent solutions often correspond to simpler activation patterns, providing a functional interpretation of coherence in terms of network behavior. This connection offers a new lens to understand generalization, establishing a link amongst solution simplicity, data alignment, and internal network representations.

To ground our theory in practice, we conduct experiments that:

1. Show coherence improves generalization and guarantee the stability of the solution. The new connection between data and optimization process that reveal the role of data in optimization.

2. Show that SAM leads to simpler (lower-rank) feature representations, indicating that coherence is reflected in the activation patterns of neural networks. This insight leverages the layered structure of deep networks and bridges theoretical analysis with the internal behavior of practical architectures.

3. Extend beyond toy settings to larger models and datasets like CIFAR-10 (see response for Response for Reivewer k5uF). Overall, our work contributes a new theoretical tool based on coherence and local stability to explain and potentially design better training algorithms that align with the structure of the data.

1. Reliance on Local Linearization: The reviewer is concerned that the paper’s theoretical analysis is limited to a linear approximation near local minima, which may not capture the full dynamics of training. They question the broader applicability of the results, given that real training often explores non-linear regions far from any critical point.

We thank the reviewer for highlighting this limitation, which we acknowledge in the main text. While our analysis is local, this regime remains highly relevant: empirical studies show that loss landscapes are often locally approximately quadratic, even in overparameterized models [1], motivating many theoretical frameworks for optimization and generalization [2, 3]. Since training often proceeds through plateau phases near minima, local dynamics are central to understanding stability. Like gradient flow and neural tangent kernel analyses, our work adopts simplifying assumptions that, while not universally valid, yield valuable insights. Our approach aligns with prior work (e.g., Dexter et al., 2024; Wu et al., 2022) demonstrating the utility of local approximations, and we will clarify these motivations in the revision.

[1] Hao Li., Visualizing the loss landscape of neural nets. NeurIPS,2018

[2] El Mehdi Achour., The loss landscape of deep linear neural networks: a second-order analysis. JMLR 2024.

[3] Kaiyue Wen. How does sharpness-aware minimization minimize sharpness? arXiv 2022.

2. Strong Assumptions and Simplifications: The reviewer notes that the analysis relies on strong simplifying assumptions, such as i.i.d. data, two-layer ReLU networks, and basic optimizers (SGD, SAM). They question whether the proposed framework of generalizing solutions captures the complexity of real neural networks, and whether the conclusions extend to other training settings like momentum or adaptive methods.

First, the i.i.d. sampling assumption is standard in theoretical deep learning (e.g., generalization bounds, stochastic processes) and reflects how minibatches are typically drawn in practice. Studying non-i.i.d. settings like curriculum or imbalanced sampling is a valuable direction for future work.

Second, while $({C}, r)$-generalizing solutions are restricted, they reflect key structures observed in deep networks—especially in activation patterns. ReLU networks exhibit rich symmetries (e.g., permutation, rescaling) that yield exponentially many equivalent solutions with identical loss, gradient, and Hessian, making our framework practically relevant despite its constraints. Such simplifications are common to enable tractable analysis.

Lastly, we focus on plain SGD and SAM to isolate the roles of noise and sharpness in a clean setting. Extending to optimizers like momentum SGD or Adam, which interact with curvature in subtle ways, is a promising direction we plan to pursue. We appreciate the reviewer’s suggestion and will note this in the revision.

3. Unclear Direct Practical Implications: The reviewer raises concern about the practicality of the proposed coherence measure, noting that it is computationally expensive and not typically used in training. They question how feasible it is to apply this concept in real-world settings as a diagnostic or design tool.

Our work is intended as a first step in highlighting coherence as a potentially valuable conceptual lens for understanding training dynamics and generalization. While the coherence measure is not yet a standard diagnostic tool, this is also true for many theoretical quantities when first introduced. A relevant example is sharpness (e.g., maximum eigenvalue of the Hessian), which was originally difficult to compute at scale, yet it inspired successful optimization strategies such as SAM that approximate the principle indirectly. Similarly, although computing exact gradient coherence or per-sample Hessian quantities is currently expensive, we see this as a motivation for future work on scalable surrogates or proxies. Our contribution is to show that coherence connects to stability in a theoretically grounded way, suggesting that it could inform the design of future optimizers or monitoring tools, even if not computed directly. We will be glad to include this perspective and the challenges of practical deployment in the discussion section of the updated version.

4. Limited Empirical Scope : The reviewer questions the empirical scope of the paper, noting that experiments are limited to a small synthetic setup. Asking whether the theoretical insights, such as coherence dynamics and feature rank reduction by SAM, hold for large-scale models and datasets, and suggest extending the experiments to larger architectures to verify applicability.

We appreciate the reviewer’s suggestion. To investigate the scalability and practical relevance of our insights, we conducted additional experiments on the CIFAR-10 dataset using a ResNet-18 model (11.7M parameters). To approximate the coherence measure in a tractable way, we used the pairwise dot product of per-sample gradients normalized by the maximum gradient norm: $\frac{\nabla l_i \nabla_jl}{\max_k ||\nabla l_k||}$ For further approximation, We compute this on a fixed subset of 100 samples (10 per class) to construct a $100 \times 100$ coherence matrix. For the feature rank calculation, we record the features from before the last linear layer for whole training dataset and use PCA to calculate the rank of feature with explain ratio up to 99.9 percent. According to our experiments, we find that the feature rank decrease as expected and the the approximated coherence measure also closely follow the training process despite milder trends due to the approximation.

5. Highly depends on several simplification assumptions: The analysis relies on simplifying assumptions (e.g., local linearization, two-layer ReLU models, specific generalizing solutions). The reviewer asks whether these assumptions are reasonable in modern training settings, and whether the results could extend to deeper architectures like CNNs or other modern structures. They also raise the question of potential new challenges in applying coherence-based stability analysis to such models (e.g., multiple dominant Hessian directions, non-linear layer interactions).?

We appreciate the reviewer’s question. While our analysis relies on standard simplifying assumptions—local linearization and two-layer ReLU networks—these are common in the literature and have yielded useful insights into deep learning dynamics (e.g., [1][2]). Despite the simplified setting, our framework captures gradient coherence and stability, which we find to remain meaningful even in large-scale models like ResNet-18 (see CIFAR-10 results). Extending this theory to deeper architectures is a promising next step. Though deeper models introduce challenges like multi-modal Hessians and inter-layer interactions, approximations such as blockwise or layer-wise coherence may offer practical paths forward, which we plan to explore.

[1] Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. COLT, PMLR, 2020.

[2] Lei Wu. Towards understanding generalization of deep learning: Perspective of loss landscapes. arXiv, 2017.

6. Measuring and Using Coherence in Practice: The paper proposes data coherence as a central concept. The reviewer asks whether this coherence measure can be feasibly computed or approximated at scale, and whether it could realistically serve as a diagnostic tool in practice.

We appreciate the reviewer’s question. We agree that coherence is a promising quantity with the potential to provide deeper insight into training dynamics and generalization. While computing the full coherence matrix exactly is impractical for large-scale models, similar to how Hessian-related quantities (e.g., sharpness or trace) are often approximated, we believe coherence can be approximated efficiently using small batches or low-rank projections. As demonstrated in our experiments (e.g., ResNet18 on CIFAR-10), such approximations can already reveal meaningful trends. Developing principled and scalable approximations of coherence, and understanding how it correlates with training outcomes, is a promising direction for future work.

7. Potential direction for design of algorithm: The reviewer asks whether the paper offers any concrete ideas or intuitions for designing new optimizers that align algorithmic behavior with data geometry such as adjusting learning rates or batch selection based on coherence, or adding regularization terms tied to per sample gradient alignment.

We thank the reviewer for the insightful suggestion. Our analysis aims to inspire optimizers that leverage gradient coherence to align with data geometry. One concrete idea is to adapt the learning rate based on coherence between mini-batches—reducing it when gradients are aligned, and increasing it when they diverge. Other possibilities include regularizers that promote gradient alignment or batch selection strategies favoring coherence. While we have not explored these experimentally, we view them as promising directions for future work.

1. Weaknesses: First, I believe related work should appear in the main body of the paper, not in an appendix. Next, all of the analysis is based on behavior near local minima, and hence the linear approximation of the loss landscape is easier to handle. Finally, the paper only studies two-layer ReLU networks. In terms of novelty, the work appears to be bordering on being incremental, but I feel there is sufficient novelty in the SAM and the 2-layer results to warrant acceptance.

We thank the reviewer for the thoughtful comments and for recognizing the novelty in our results. We have tried to provide all contextualized related works and cited them inline in the introduction itself. As such, we did not feel a separate related works section (which is mostly 'additiona related works') added a lot of value over what is already in the intro. If the reviewer strongly feels otherwise, we are happy to move it to the main paper. We agree that it plays an important role in framing the contributions of the paper.

Regarding novelty, our coherence-based analysis offers a new perspective on how property of data interacts with optimization. Unlike prior works that focus solely on sharpness, we show that alignment in the loss landscape—driven by data coherence—can act as a stabilizing factor, revealing a previously underexplored mechanism for generalization. This insight leads to a theoretical breakthrough by integrating data-dependent property into the analysis of optimization and generalization. Furthermore, our framework bridges loss landscape geometry, distribution of data, and simplicity bias in neural networks. We view this as a key advancement, as simplicity bias is a defining feature of modern deep networks and often touted as an important aspect of explaining generalization (see ref[1][2]). Our results thus help connect optimization theory with broader aspects of deep learning, filling a critical gap in existing understanding.

On the focus on local analysis, we fully acknowledge this as a limitation. Any approximation of neural networks for theoretical tractability, including neural tangent kernels, infinite widths of intermediate layers, gradient flows, bayesian perspective on sgd, make some assumptions that may not always hold. However, we believe the characterization of behavior around local minima is still highly relevant in modern deep learning and is grounded in empirical observations. Training dynamics often enter a hovering regime where the loss plateaus and the optimizer oscillates near a minimum for a substantial portion of training. This regime is known to influence generalization and stability properties. Our analysis is intended as a first step to provide a rigorous foundation for understanding this important phase of training. Lastly, while it is true that our theoretical analysis is grounded in two-layer ReLU networks, we emphasize that this model class has been widely used to uncover foundational insights in deep learning theory. (see ref [3][4]). The simplifications being made for RELU constructions are highly relevant to modern data and neural network interactions. Moreover, we extend our experimental validation to deeper architectures (e.g., ResNet-18 on CIFAR-10) and show the correspondence to coherence and feature-rank trends unearthed through our theoretical analysis, suggesting that the insights transfer beyond the two-layer setting. We agree that generalizing the theory to deeper or more complex networks is an exciting and important future direction, but we hope the reviewer sees the importance of our work as a stepping stone in that direction. We appreciate the reviewer’s overall positive assessment and constructive feedback.

[1] K. Gatmiry, Z. Li, S. J. Reddi, and S. Jegelka, “Simplicity bias via global convergence of sharpness minimization,” arXiv, 2024

[2] Springer, J. M., Nagarajan, V., and Raghunathan, A. (2024). Sharpness-aware minimization enhances feature quality via balanced learning. ICLR.

[3] El Mehdi Achour., The loss landscape of deep linear neural networks: a second-order analysis. JMLR 2024.

[4] Kaiyue Wen. How does sharpness-aware minimization minimize sharpness? arXiv 2022.

2. I found Definition 2 of -generalizing solution a bit ill-defined and unnatural. What is a (C,r) generalizing solution? You define and , but do not state what a generalizing solution is.

We thank the reviewer for pointing this out and apologize for the lack of clarity. In our work, a generalizing solution refers to a solution (i.e., a set of network weights) that achieves zero test classification error under the data distribution we mentioned in the main context (section 3.3). (i.e., $y = f_w(x) \;\;\forall x,y \sim \mathcal{D}$) To be more specific, we provide one example in the appendix C.1. To provide intuitive understanding, the C in (C,r) solution controls how many neurons are activated in the first layer, and r control the the flatness in the solution as mentioned in the Definition 2 and the paragraph underneath. We will revise the main text to more clearly introduce what we mean by a generalizing solution before presenting Definition 2. If the reviewer has a specific concern about the formulation, we would be happy to address it in greater detail.

Note that the generalizing solutions encompass exponentially many parameter configurations due to inherent symmetries in neural networks such as hidden unit permutations and layer-wise scaling invariance. These transformations preserve important properties such as loss, gradient, and Hessian, resulting in a broad equivalence class of solutions with indistinguishable local geometry. Our work provides a principled characterization of the optimization dynamics within this rich class. Moreover, it demonstrates the compatibility of our framework with broader aspects of deep learning, including simplicity bias and representation learning.

1. The novelty of this paper with respect to other work on dynamical stability (such as Wu & Su (2023)) is unclear. In Appendix B, the authors attempt to distinguish their work as, "... but compared to prior works [Wu et al., 2018, 2022, Wu & Su, 2023], instead of assuming mean square loss, we have more general abstraction to include different kind of loss function." This is not true: for example, Wu & Su (2023) use a generic loss function and Hessian throughout their analysis as well (similar to this paper), and only use the MSE as an example in Section 2. The rest of their work studies a general ERM problem, with a gradient and a Hessian (though their paper uses the associate empirical Fisher matrix).

We thank the reviewer for the helpful clarification. Upon revisiting Wu & Su (2023), we agree that their broader framework is presented in terms of general empirical risk minimization, and that their analysis of generalization error and sharpness indeed applies to a general loss and Hessian (or Fisher). Our statement in Appendix B was too strong and will be revised accordingly. That said, we would like to clarify a subtle but important distinction: while their overall setup is general, the key linear stability analysis in Wu & Su (2023) especially the parts relying on spectral properties of the Jacobian or stability matrix is carried out under the mean squared error (MSE) loss. The use of MSE enables concrete expressions for the Hessian and simplifications that arise from its symmetry and quadratic form. For instance, Proposition 3.2 builds on Equation (3), which assumes MSE, leading to the specific Hessian form in Equation (2). Proposition 3.4 similarly uses this structure, and its derivation in Equation (23), as noted in Appendix B.2, depends on Lemma 3.3 — both of which are MSE-specific. While Wu & Su later discuss generalization using broader loss functions, their linear stability bounds are derived concretely under the MSE setting. We hope this clarifies the intended distinction, and we appreciate the opportunity to correct and clarify our prior statement.

2. This paper would benefit from a brief explanation of Dexter et al. (2024), since that work is the basis for some of the theory and the experimental setup. As such, it is useful to readers to understand why the specific experimental setup is chosen.

We appreciate the reviewer for the question. Our work generalizes the stability framework of Dexter et al 2024 on SGD analysis to include Sharpness-Aware Minimization (SAM) and random perturbation SGD, revealing that SAM imposes a stricter and more geometry-sensitive stability condition than SGD. It proves that SAM preferentially selects solutions with higher inter-sample curvature alignment, even when sharpness is held constant. Unlike Dexter et al, which focuses solely on SGD, this paper formalizes the link between stability and simplicity bias across optimizers and instantiates the theory in two-layer ReLU networks, distinguishing generalizing from memorizing solutions by their coherence structure.

3. After Theorem 3.1, the paper states, "Further, in the stable regime, the iterates do not converge exactly to the minimum, but instead remain in a bounded region around it (part 3)". Is this true for all minima, or only when w*=0? To me, part 3 seems to instead suggest that while the iterates remain in a bounded region, that region may not be symmetric around the minimum; is this accurate?

We thank the reviewer for the thoughtful question. We clarify that the statement following Theorem 3.1 holds more generally and is not restricted to the case where w*=0. Since we are interested in dyanamics around optima actual value of $w^*$ does not matter. The key point is that due to the presence of additive noise (independent of the curvature or symmetry of the loss landscape) iterates do not converge exactly to the minimizer, but instead fluctuate within a bounded region. If we look at equation (3), we can find that even under w*=0, there will still be non-zero update of the weight. This is not the result of the loss landscape but the noise that is added into the algorithm artificially. Here, we provide several random perturbation base algorithms (see ref [1][2]) that gradually decrease the noise during the training to achieve better convergence behavior as support for our statement.

[1] Bisla., A. Low-pass filtering sgd for recovering flat optima in the deep learning optimization landscape. AISTATS, 2022.

[2] Yong Liu. Random sharpness-aware minimization. NeurIPS, 2022.

4. In Definition 2, it would be beneficial to use a consistent notation.

We thank the reviewer for addressing the issue. We will modify it accordingly in the updated version.

5. In Appendix C.1, $e_1$ is not defined.

We appreciate reviewer for pointing that out. $e_1$ is the Canonical orthonormal basis with 1 as its first element and zero otherwise. We will incorporate this definition in the updated version.

6. Could you better contextualize the novelty of this paper and distinguish it from Wu Su (2023)? I'm not convinced by the appendix's statement that the MSE loss being used is a limitation of that work that is addressed here.

We see our novelty in two different parts compared to prior work:

First, Coherence base analysis: We use the Hessian coherence matrix as a central object to capture alignment between per-sample gradients. This sheds light on how SGD and SAM interact differently with different data distributions — particularly in terms of stability and generalization — an aspect that is not a focus of Wu & Su (2023). In contrast to their emphasis on the noise structure interaction with Fisher information matrix, our analysis reveals how data geometry influences the stability of solutions through coherence. This directly provides a theoretical bases for empirical observations on simplicity bias made by other papers. (see [1] [2])

Second, bridging linear stability to neural network structure via activation patterns: A key novelty of our work is connecting linear stability theory to the activation patterns in ReLU networks. We show that generalizing solutions tend to induce more coherent gradients, which in turn correspond to simpler (i.e. more sparse, and more structured activation patterns. This provides a concrete way to interpret algorithmic bias (e.g., from SAM) in terms of neural architecture behavior — a perspective not developed in Wu & Su or related stability works. We hope these clarifications help highlight how our work complements and extends the existing body of literature on optimization dynamics.

[1] K. Gatmiry, Z. Li, S. J. Reddi, and S. Jegelka, “Simplicity bias via global convergence of sharpness minimization,” arXiv, 2024

[2] Springer, J. M., Nagarajan, V., and Raghunathan, A. (2024). Sharpness-aware minimization enhances feature quality via balanced learning. ICLR.

Thank you for the clear rebuttal. The authors have addressed my primary concern (novelty), citing specific parts of prior work that makes the assumptions that their paper claims. The rebuttal also clearly addresses the other issues I raised. I currently do not see any reason to reject this paper, so I will update my score to a 5.