PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Understanding Sharpness-Aware Minimization

Maksym Andriushchenko, Nicolas Flammarion

arXiv:2206.06232v1cs.LG

TL;DR

The paper asks why SAM and especially low-m sharpness improve generalization when existing PAC-Bayes and flat-minima explanations are incomplete. It analyzes SAM’s implicit bias and stochastic convergence theoretically, studies non-linear networks empirically, and finds that low-m objectives and SAM-selected solutions can support better generalization.

  • Problem

    Existing PAC-Bayes and flat-minima explanations do not fully explain SAM’s generalization gains or the benefit of low-m sharpness.

  • Method

    The paper analyzes SAM’s implicit bias for diagonal linear networks, studies non-linear networks empirically, and proves stochastic convergence results for non-convex objectives.

  • Results

    SAM’s implicit bias is related to the ℓ1-norm of the resulting linear predictor, is stronger for 1-SAM than n-SAM, and fine-tuning ERM models with SAM can improve generalization.

  • Takeaways & Limitations

    Low-m SAM’s generalization benefit likely comes from its better objective, while convergence analysis connects the method’s optimization behavior with its empirical performance.

  • Takeaways & Limitations

    The existing generalization bound cannot explain the improvement of low-m SAM, and convergence to flatter minima does not always imply better generalization.

Abstract

from arXiv · show

Sharpness-Aware Minimization (SAM) is a recent training method that relies on worst-case weight perturbations which significantly improves generalization in various settings. We argue that the existing justifications for the success of SAM which are based on a PAC-Bayes generalization bound and the idea of convergence to flat minima are incomplete. Moreover, there are no explanations for the success of using $m$-sharpness in SAM which has been shown as essential for generalization. To better understand this aspect of SAM, we theoretically analyze its implicit bias for diagonal linear networks. We prove that SAM always chooses a solution that enjoys better generalization properties than standard gradient descent for a certain class of problems, and this effect is amplified by using $m$-sharpness. We further study the properties of the implicit bias on non-linear networks empirically, where we show that fine-tuning a standard model with SAM can lead to significant generalization improvements. Finally, we provide convergence results of SAM for non-convex objectives when used with stochastic gradients. We illustrate these results empirically for deep networks and discuss their relation to the generalization behavior of SAM. The code of our experiments is available at https://github.com/tml-epfl/understanding-sam.

1. Introduction

The paper argues that existing explanations for SAM’s generalization gains are incomplete and investigates its implicit bias and convergence behavior theoretically and empirically.

  • SAM: SAM applies worst-case weight perturbations during training to minimize a sharpness-related objective.The main version considered uses one gradient-ascent step to approximate the perturbation problem.
  • Motivation: The paper studies why SAM generalizes better, challenging explanations based on PAC-Bayes bounds and convergence to flat minima.It identifies unresolved questions about worst-case perturbations and the role of low-m m-sharpness.
  • Contributions: The paper tests why maximizing over fewer training points can improve generalization and concludes that the benefit likely comes from a better objective.This addresses the unexplained advantage of low-m SAM.
  • Contributions: The authors theoretically analyze SAM’s implicit bias for diagonal linear networks and empirically relate it to non-linear networks.The study focuses on how the optimization algorithm selects among solutions fitting the training data.
  • Contributions: The authors prove stochastic convergence results for SAM on non-convex objectives and check convergence empirically for deep networks.They relate these optimization results to SAM’s generalization behavior.

2. Background on SAM

This section defines sharpness and SAM variants, explains the role of batch size and worst-case perturbations, and summarizes practical implementation choices and motivating experiments.

  • Sharpness: Sharpness measures how rapidly a classifier’s loss changes within a parameter neighborhood on a selected set of training points.The paper defines sharpness for any subset S of the training set.
  • Sharpness: m-sharpness averages sharpness over all training batches of size m, whereas traditional sharpness uses the full training set.Thus m controls the subset size used in the sharpness calculation.
  • SAM variants: n-SAM maximizes the sum of losses over all training points, while m-SAM maximizes over batches of m points and is related to m-sharpness.Foret et al. theoretically motivate n-SAM, although m-SAM is reported to generalize better empirically.
  • Optimization: SAM is implemented with stochastic gradients using the same batch for the inner perturbation step and outer gradient step.Gradient normalization can be included, but the paper omits it theoretically because it is not necessary for improving generalization.
  • Empirical motivation: Only low-m SAM substantially improves generalization in the compared perturbation methods, while random perturbations provide marginal improvement.The experiment compares ERM, random perturbations, n-SAM, and 128-SAM on CIFAR-10 and CIFAR-100 models.

3. Challenging the Existing Understanding of SAM

The section argues that existing PAC-Bayesian and flat-minimum explanations do not account for SAM’s generalization behavior, while m-sharpness is not universally predictive. Empirical analyses expose these limitations and motivate a more nuanced understanding of SAM.

  • Limitations of flat-minimum explanations: Convergence to flatter minima measured by m-sharpness does not always imply better generalization, so sharpness is not universally decisive.The section motivates studying other quantities that may correlate with m-sharpness only in some cases.
  • Limitations of existing generalization bounds: The PAC-Bayesian justification analyzes random parameter perturbations rather than SAM’s worst-case perturbations.Random perturbations and n-SAM do not substantially improve generalization in the reported empirical comparison.
  • Limitations of existing generalization bounds: Existing SAM generalization bounds incorrectly rank 128-SAM as worst despite it being the only method that improves generalization.The bound extension to m-SAM would therefore fail to explain low-m SAM’s observed benefit.
  • Empirical examination of m-sharpness: All evaluated group-normalized models achieve zero training error while their m = 128 sharpness is compared across perturbation radii, batch sizes, and ERM versus SAM.The models use large batches of 1024 or small batches of 128.
  • Empirical examination of m-sharpness: The 1-sharpness of a linear model is independent of the parameter vector, making all global training-loss minimizers equally sharp under this measure.This provides a simple example where sharpness cannot distinguish minima with different generalization behavior.

4. Understanding the Generalization Benefits of SAM

The paper argues that low-m SAM’s generalization benefit is better explained by its implicit bias than by inner-maximization accuracy or batch-normalization regularization. Theory on diagonal linear networks and experiments on nonlinear networks link SAM’s behavior to sparsity, training dynamics, and fine-tuning.

  • Testing Two Natural Hypotheses: Lower m’s improvement is not specific to batch normalization, motivating an implicit-bias explanation for low-m SAM.The authors test whether lower m improves maximization accuracy or batch-normalization regularization before favoring the implicit-bias hypothesis.
  • Implicit Bias in Diagonal Linear Networks: 1-SAM and n-SAM bias full-batch dynamics toward φα-minimizing solutions with effective scales below α, and 1-SAM typically has the stronger sparsity-inducing bias.The theory states that ||α_1-SAM||_1 < ||α_n-SAM||_1 typically, making 1-SAM closer to the sparse regime.
  • Implicit Bias in Diagonal Linear Networks: On sparse regression, 1-SAM generalizes better than n-SAM and ERM, while ERM and n-SAM have approximately similar performance.The comparison remains favorable to 1-SAM when the perturbation radius is tuned separately; n-SAM then provides only marginal improvements.
  • Empirical Study of Nonlinear Networks: SAM favors simpler, more stable ReLU interpolations and appears biased toward sparse combinations of ReLUs across five random initializations.The observed behavior is related to sparsity in the resulting linear predictor and to low-dimensional structures in ReLU models.
  • Empirical Study of Nonlinear Networks: Switching methods shows that early training has little influence on final performance, whereas enabling SAM late can still significantly improve generalization.The authors interpret this as evidence against convergence to entirely distinct minima and report that ERM-to-SAM fine-tuning can substantially improve a pretrained model.

5. Understanding the Optimization Aspects of SAM

The paper fills a convergence gap for stochastic SAM and relates its training-loss convergence to empirical behavior in deep networks, including noisy-label settings where convergence can promote overfitting.

  • Theoretical convergence: SAM convergence analysis had been missing conditions guaranteeing convergence or showing convergence to a stationary point of the ERM objective.The analysis targets m-SAM under assumptions on bounded variance, individual smoothness, and, for stronger guarantees, the Polyak-Lojasiewicz condition.
  • Theoretical convergence: SAM recovers SGD convergence rates under the usual condition on the descent step size.
  • Theoretical convergence: The ascent step size ρt may decrease as O(√γt), more slowly than the descent step size, while still ensuring convergence.The authors relate this condition to observations from neural-network training.
  • Theoretical convergence: The convergence proof uses alignment between the SAM step and the ordinary gradient step.The stated bound lower-bounds their inner product by (1 − ηβ)∥∇L(wt)∥2.
  • Deep-network experiments: On CIFAR-10, ERM and SAM reach nearly zero training loss, with 0.0013 ± 0.00002 for ERM versus 0.0034 ± 0.0004 for SAM.The corresponding CIFAR-100 behavior is qualitatively similar.
  • Deep-network experiments: With 60% label noise, SAM improves generalization over ERM but later fits noisy points, so early stopping remains necessary.The authors also observe that SAM’s benefit appears along the optimization trajectory rather than only near a minimum.

6. Conclusions

The paper concludes that SAM’s generalization cannot be fully explained by existing sharpness-based justifications and instead links its behavior to implicit bias and convergence analysis.

  • Conclusions: Existing explanations of m-SAM based on generalization bounds and convergence to flat minima are incomplete.
  • Conclusions: For diagonal linear networks, 1-SAM and n-SAM implicitly minimize quantities related to the ℓ1-norm of the resulting predictor, with a stronger effect for 1-SAM.
  • Conclusions: Fine-tuning an ERM model with SAM can produce significant generalization improvements on non-linear networks.
  • Conclusions: The paper provides stochastic-gradient convergence results for non-convex objectives and confirms them empirically on deep networks.

A. Implementations of the SAM Algorithm in the Full-Batch Setting

This section defines full-batch implementations of 1-SAM and n-SAM, relates them to gradient descent at perturbed intermediate points, and introduces the diagonal-linear setting used to study implicit bias.

  • Full-batch implementations: The full-batch implementations analyze the two extreme choices m = 1 and m = n through their corresponding objectives.
  • Full-batch implementations: Their updates use gradient descent with gradients evaluated at intermediate points, whose construction differs between 1-SAM and n-SAM.
  • Full-batch implementations: The practical stochastic-gradient version of SAM corresponds to the m-SAM algorithm defined in the main text.
  • Implicit-bias setup: The implicit-bias analysis uses noiseless regression with a two-layer diagonal linear network, where different optimization algorithms can be characterized precisely.
  • Implicit-bias setup: Gradient flow from initialization scale α converges to the minimum ℓ1-norm interpolator for small α and the minimum ℓ2-norm interpolator for large α.The interpolating regularizer φα transitions between these norms as the initialization scale changes.

B.1. Implicit Bias of the n-SAM Algorithm.

The implicit-bias analysis characterizes n-SAM and 1-SAM through a potential φα and shows how perturbation size modifies the effective initialization scale governing the selected interpolator.

  • n-SAM implicit bias: n-SAM converges, under the stated interpolation condition, to a minimum-norm solution characterized by the KKT conditions for φα.
  • n-SAM implicit bias: For sufficiently small ρ, n-SAM has an effective bias parameter smaller than the initialization scale α.The associated bias depends on an integral involving the training trajectory and is typically negligible because of an additional d/n factor.
  • n-SAM implicit bias: The n-SAM derivation represents the perturbed iterate as wsam(t) = w(t) + ρ∇L(w(t)) and analyzes its residual dynamics.
  • 1-SAM implicit bias: For 1-SAM, bounded inputs and bounded trajectories imply α1-SAM,i ≤ αi for every coordinate when ρ ≤ 1/[4R^2√(B(B+∥β∗∥2))].

B.3. Comparison between 1-SAM and n-SAM

The analysis compares 1-SAM and n-SAM through their implicit bias, convergence behavior, and switching experiments. It finds stronger biasing and generalization effects for 1-SAM, while both methods retain convergence guarantees under stated smoothness conditions.

  • Theoretical comparison: In the overparametrized regime d >> n, 1-SAM’s biasing effect is typically O(n) times stronger than n-SAM’s.The comparison uses the ℓ1-norms of trajectory-dependent quantities associated with the two methods.
  • Empirical comparison: 1-SAM generalizes significantly better than n-SAM even when each method uses its optimally tuned perturbation radius.Under a grid search over ρ, n-SAM provides only marginal improvements while 1-SAM still outperforms both ERM and n-SAM.
  • Implicit bias: The implicit bias of SAM depends on the integral of ρ_s times the loss along the optimization trajectory, which lowers the effective initialization scale and favors sparser solutions.The resulting effective scale is α_1-SAM-ρs = αe^−ρ∆_1-SAM-ρs+O(ρ^2).
  • Switching methods: Switching from ERM to 1-SAM can improve generalization because the post-switch trajectory still accumulates sufficient biasing effect.The loss interpolation indicates that the resulting minima are linearly connected and lie in the same basin.
  • Switching methods: Switching from 1-SAM back to ERM preserves the implicit bias accumulated during the initial SAM phase and does not negatively affect generalization.The first 1000 epochs with SAM can make the effective initialization scale low enough for the bias to persist.
  • Convergence: For smooth non-convex objectives, SAM converges under γ < 1/β and ρ < 1/β, recovering gradient-descent rates with constants that increase with ρ.Stronger results hold under the Polyak–Łojasiewicz condition, while convex objectives permit a broader step-size condition.

C.2.1. CONVERGENCE OF n-SAM

The paper establishes convergence results for stochastic n-SAM by controlling the expected update direction and objective decrease under smoothness and bounded-variance assumptions. The resulting bounds include diminishing-step-size guarantees and stronger rates under additional conditions.

  • Algorithm: Stochastic n-SAM uses different mini-batches for its ascent and descent steps.This distinguishes the analyzed n-SAM update from m-SAM, which reuses the same batch.
  • Proof strategy: The SAM update is shown to remain sufficiently aligned with the population gradient, enabling objective-decrease bounds for stochastic updates.The analysis uses auxiliary lemmas for update alignment and function-value decrease.
  • Proof strategy: The convergence analysis controls stochastic n-SAM similarly to stochastic gradient descent under bounded variance and smoothness assumptions.The proof compares the noisy update with a noiseless gradient step and bounds the stochastic remainder using the mini-batch variance.
  • m-SAM comparison: The m-SAM convergence proof follows the same overall structure but requires Lipschitz individual gradients to control expected update alignment.The same-batch ascent and descent structure is treated separately from the different-batch n-SAM analysis.
  • Convergence guarantees: With diminishing step sizes γ_t = 1/√T and ρ_t = 1/(T^1/4β), the analysis obtains the same convergence result as the corresponding theorem up to simplified constants.The result is derived from the more general theorem for the stochastic SAM variant.

D. Experimental Details

The experiments evaluate SAM across datasets, architectures, perturbation radii, batch sizes, model widths, and normalization choices. Results consistently emphasize the roles of smaller m, batch size, intermediate width, and normalization-aware sharpness measurement.

  • Training setup: Experiments use SGD with momentum and weight decay on CIFAR-10 and CIFAR-100, primarily with pre-activation ResNets and batch size 128.Basic random-crop and mirroring augmentation is applied throughout.
  • Training setup: The perturbation radius ρ is selected by grid search, with ρ = 0.1 optimal in most experiments.The ERM → SAM experiment instead uses ρ = 0.4 for CIFAR-10 and ρ = 0.2 for CIFAR-100.
  • Effect of m: Smaller m produces larger generalization improvements in m-SAM, with the improvement varying continuously with m.These results use group normalization and fixed batch size 256.
  • Effect of batch size: When m equals batch size, excessively large m yields only marginal generalization improvements of approximately 0.2%.It does not close the gap between large-batch 1024 SGD and small-batch 256 or 128 SGD.
  • Effect of model width: SAM’s generalization improvement is largest at intermediate model widths, such as width factor 16, rather than at the widest models.Small networks show little or no improvement because additional regularization may not be needed.
  • Sharpness measurement: BatchNorm makes sharpness measurements differ substantially between training-time and test-time behavior, so the main results use GroupNorm.The discrepancy arises because SAM uses training-time BatchNorm while post-hoc sharpness commonly uses test-time BatchNorm.

E.5. Training Loss for ERM vs. SAM Models

Additional experiments compare ERM and SAM training objectives and examine changing perturbation radii. The objective decreases for both methods, while decreasing ρ over training is reported as detrimental to generalization.

  • Training objective: Training and test errors slightly increase after the first learning-rate decay, despite the total training objective continuing to decrease.This indicates that the error increase is not attributed to failure to optimize the stated objective.
  • Training objective: The total ERM and SAM training objectives both show a clear decreasing trend over epochs.The objective includes the ℓ2 regularization term.
  • Perturbation schedule: Decreasing the perturbation radius ρ_t according to the learning-rate schedule is detrimental to generalization compared with constant-radius SAM.Both variants use ℓ2-normalized updates.

E.7. Experiments with Noisy Labels

Experiments on CIFAR-10 and CIFAR-100 with 60% noisy labels compare SAM with and without gradient normalization. After selecting a fixed inner step size, unnormalized SAM nearly matches standard SAM in best test error.

  • Experiments use CIFAR-10 and CIFAR-100 with 60% noisy labels to compare SAM variants.
  • The comparison evaluates SAM with a fixed inner step size ρ without gradient normalization, called unnormalized SAM.The fixed ρ was selected through a prior grid search.
  • The best test error over epochs for unnormalized SAM almost exactly matches standard SAM.
Loading 2206.06232v1…