Sharpness-Aware Gradient Descent (SA-GD)

Updated 9 November 2025

SA-GD is a family of first-order optimization methods that perturbs model weights to explicitly account for the sharpness of the loss landscape.
It employs a bi-level objective with an inner maximization (to find worst-case perturbations) and an outer descent step, effectively guiding models toward flatter, more generalizable minima.
Variants like SAM, GCSAM, mSAM, and others enhance efficiency and stability through strategies such as gradient centralization, micro-batch averaging, and renormalization.

Sharpness-Aware Gradient Descent (SA-GD) is a family of first-order optimization algorithms for neural network training that enforce robustness to worst-case perturbations of model weights. These methods augment empirical risk minimization by explicitly accounting for the “sharpness” of the local loss landscape, encouraging convergence to flatter regions that are hypothesized to generalize better in over-parameterized models. The class subsumes Sharpness-Aware Minimization (SAM) and its numerous variants, and has seen widespread empirical success as well as intensive theoretical scrutiny.

1. Mathematical Formulation and Core Algorithm

The central object of SA-GD is the bi-level sharpness-aware empirical objective: $\min_{w} \max_{\|\epsilon\|_{p} \le \rho} L_S(w + \epsilon)$ where $L_S$ is the loss over a mini-batch $S$ , $\rho > 0$ the “sharpness radius,” and $\|\cdot\|_p$ denotes a norm constraint (typically $p=2$ ). This inner maximization evaluates the worst-case loss climb within a local ball, making $w$ robust to small weight perturbations.

In practice, the inner maximization is approximated by a single first-order Taylor expansion: $\epsilon^* = \rho \frac{\nabla L_S(w)}{\|\nabla L_S(w)\|_q}$ where $1/p + 1/q = 1$. For $p=2$ , $q=2$ .

After finding $\epsilon^*$ , the outer minimization is performed via a descent step at the perturbed point: $w \leftarrow w - \alpha \nabla L_S(w + \epsilon^*)$ with learning rate $\alpha$ .

Pseudocode: Vanilla SAM/SA-GD

for t = 0,...,T-1:
    Sample mini-batch B
    g = ∇L_B(w)
    ε = ρ * g / ||g||_2
    g' = ∇L_B(w + ε)
    w = w - α * g'

This structure doubles the per-step computation relative to vanilla SGD due to the requirement for two forward/backward passes.

2. Theoretical Properties, Landscape Dynamics, and Flatness

Sharpness-aware steps bias the optimization trajectory toward “flat” minima—regions where the loss is locally insensitive to parameter perturbations. For convex quadratics, SAM/SA-GD converges to a two-point limit cycle oscillating across the minimum in the direction of maximal curvature, but in non-quadratic regimes, the oscillatory component is superposed with a slow drift that acts to minimize the spectral norm of the Hessian: $w_{t+1} - w_t = -2\epsilon_t - \frac{\alpha \rho^2}{2} \nabla \lambda_{\max}(H) + O(\alpha \rho^3)$ Thus, beyond flattening by design, SA-GD implements implicit regularization that pulls iterates toward “wide minima”—low Hessian operator norm—believed to offer better out-of-distribution generalization (Bartlett et al., 2022).

In stochastic convex optimization, SA-GD can converge rapidly to flat empirical minima; however, in settings where flat empirical minima do not generalize, population risk can remain $\Omega(1)$ even as the empirical sharpness vanishes, indicating that flatness alone is insufficient to guarantee generalization (Schliserman et al., 5 Nov 2025).

3. Variants: Efficiency, Robustness, and Limitations

A range of SA-GD derivatives have been proposed to target computational, stability, or invariance limitations:

Gradient Centralized Sharpness-Aware Minimization (GCSAM): Incorporates mean-centering of weight gradients before the ascent step, which reduces gradient variance and suppresses noise spikes. GCSAM consistently outperforms standard SAM and traditional optimizers in both accuracy and computational efficiency, incurring only negligible overhead above SAM. The mean-subtracted gradient yields a “tighter” adversarial region, with $\|\nabla_{GC}\|_2 \leq \|\nabla\|_2$ and greater stability (Hassan et al., 20 Jan 2025).
Micro-Batch-Averaged SAM (mSAM): Splits each batch into $m$ shards, computes per-shard SAM perturbations, aggregates the resulting gradients, and applies a single update. mSAM provably finds minima with even lower Hessian spectral norm than SAM and offers improved generalization on vision and NLP tasks. Computational overhead is marginal due to full parallelizability and the use of smaller micro-batch computations (Behdin et al., 2023).
K-SAM: Approximates full-batch ascent and descent gradients using only the top- $K$ hardest (highest-loss) examples per batch for each step, dramatically reducing total computation. K-SAM attains test accuracy on par with SAM, but runs at near-SGD speed and is highly amenable to large-scale/distributed training. Backpropagation is applied only on selected samples, reducing both time and memory consumption (Ni et al., 2022).
StableSAM (SSAM): Addresses potential instability in SAM (due to gradient norm amplification from the ascent step) by rescaling the descent gradient so its norm matches that of the original unperturbed gradient. This renormalization permits higher learning rates, lowers excess risk stability bounds, and achieves improved empirical generalization, with minimal implementation overhead (Tan et al., 14 Jan 2024).
Monge SAM (M-SAM): Replaces the Euclidean norm in the adversarial step with the Riemannian metric naturally induced by the loss surface (the Monge metric). This yields a reparameterization-invariant update, preventing metric-induced distortions of sharpness under parameter transformations. M-SAM interpolates between standard SAM and vanilla gradient descent: in steep regions it acts like GD, and in flat regions like SAM. It empirically attains flatter minima and outperforms SAM in challenging fine-tuning tasks, showing greater robustness to hyperparameter settings and less attraction to suboptimal equilibria (Jacobsen et al., 12 Feb 2025).

4. Generalization: Flatness, Feature Balancing, and Caveats

Traditionally, flat minima (quantified by low worst-case loss climb in parameter neighborhoods or small Hessian spectral norm) are thought to correlate with improved generalization. However, recent research challenges the sufficiency of flatness: there exist scenarios where flat empirical minima yield poor population risk (Schliserman et al., 5 Nov 2025).

An alternative mechanism for generalization improvement is feature balancing. In heterogeneous datasets, sharpness-aware updates downweight overfit features and upweight underlearned ones. This mechanism encourages models to exploit multiple predictive signals rather than collapsing onto the simplest or most spurious feature, a phenomenon evidenced by more equalized importance weights and improved feature-probe accuracy on datasets containing redundant or spurious features (e.g., CelebA, Waterbirds, CIFAR-MNIST, DomainBed). The “learning-rate scaling effect” of SAM dampens updates on well-learned features and accelerates those on poorly learned ones, leading to systematically improved representation quality (Springer et al., 30 May 2024).

5. Empirical Performance and Implementation Guidelines

Across a range of architectures and domains, SA-GD methods consistently surpass classic optimizers such as SGD and Adam in test accuracy, with especially strong gains in medical imaging and out-of-distribution generalization (Hassan et al., 20 Jan 2025). For example, on CIFAR-10 and a COVID-19 chest X-ray task, GCSAM yields higher test accuracy (+0.4–1.5% and +2–3%, respectively) while maintaining computational overhead only marginally above SAM. mSAM further reduces the largest Hessian eigenvalue and achieves incrementally higher test accuracy on CIFAR-100, ImageNet-1k, and diverse NLP benchmarks (Behdin et al., 2023). K-SAM attains SAM-level accuracy on large-scale data at less than half the wall-clock time (Ni et al., 2022). SSAM grants additional stability and accuracy across a spectrum of tasks (Tan et al., 14 Jan 2024).

Key empirical insights:

All SA-GD variants require two gradient evaluations per iteration ( $\approx 2\times$ SGD cost); efficiency-enhancing variants reduce this substantially via subset, micro-batch, or normalization tricks.
Gradient centralization, micro-batch averaging, and renormalization are effective low-cost regularizers against noise and instability.
Choice of the ascent radius $\rho$ is critical; values are typically in $[0.01, 0.1]$ gradient norm units for vision models.
Strong momentum and moderate/small batch sizes help mitigate saddle-point attraction and instability in SA-GD (Kim et al., 2023).
Hyperparameter robustness is improved in invariant (M-SAM) and stabilized (SSAM) variants.

6. Limitations, Pathologies, and Open Problems

Saddle Point Attraction: Theoretical analysis shows that for certain choices of $\rho$ , SA-GD can turn saddle points into attractors, slowing escape from high-loss plateaus. This effect can be mitigated by increasing noise via smaller batch sizes or injecting momentum, but highlights sensitivity to algorithmic tuning (Kim et al., 2023).
Flatness Is Not Sufficient: In convex settings, SA-GD can converge to flat minima with high empirical sharpness but fail to generalize, indicating that flatness, without other structural alignment with the data distribution, may not guarantee low population risk (Schliserman et al., 5 Nov 2025).
Computation: While plain SAM doubles per-step cost, variants such as K-SAM or mSAM alleviate but do not eliminate this burden. Some methods still require a full forward pass for per-sample loss computation at each step (Ni et al., 2022).
Parameterization Dependence: Euclidean-norm-based SA-GD is not invariant under parameter reparameterizations, leading to metric artifacts in sharpness; Riemannian-based methods (M-SAM) address this at little extra computational cost (Jacobsen et al., 12 Feb 2025).
Hyperparameter Sensitivity: All SA-GD methods require careful tuning of $\rho$ and step size $\eta$ ; stabilized variants (SSAM, M-SAM) are more robust but not wholly insensitive.
Generalization Mechanisms: There is ongoing debate on whether the empirical correlation of flatness and generalization is causal or merely an artifact. Alternative explanations (e.g., feature balancing, regularization via sharpness) are being explored (Springer et al., 30 May 2024).

7. Summary Table: Algorithmic Variants and Key Metrics

Variant	Extra Cost vs SGD	Flatness Control	Generalization Gains	Notable Features
SAM	$\approx\times 2$	$\ell_2$ norm, Euclidean	+0.4–1.3% (CIFAR), +2–3% (med img)	Simple, robust, 2x computation
GCSAM	$\approx\times 2.1$	GC-weighted ascent direction	Outperforms SAM across domains	Mean-centring stabilizes ascent
mSAM	$\gtrsim \times 2$	Per-microbatch adversarial	Flattest minima, highest test acc	Efficient for large-scale/parallel
K-SAM	$\sim \times 1$	Top- $K$ hardest samples	SAM-level accuracy, faster than SGD	Subset-restricted, minimal overhead
SSAM	$\sim \times 2$	Renormalized step	Marginal accuracy, broader LR regime	Step-rescaling, improved stability
M-SAM	$\sim \times 2$	Riemannian (Monge) metric	Up to +10% vs. SAM (alignment task)	Invariant to reparametrization, smoother

Throughout, careful implementation of SA-GD optimizers—together with appropriate stabilization strategies and hyperparameter tuning—remains essential for leveraging their regularization properties while minimizing computational and dynamical risks. These methods continue to inform both practical deep learning and the theoretical understanding of generalization phenomena in modern neural networks.