PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Sharpness-aware Minimization for Improved Training of Neural Networks

Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, Vincent Y. F. Tan

arXiv:2110.03141v2cs.AIcs.CVcs.LG

TL;DR

Overparameterized DNNs can generalize poorly, while SAM improves generalization at roughly double the computational cost of base optimizers. ESAM addresses this trade-off with stochastic weight perturbation and sharpness-sensitive data selection, achieving lower overhead with preserved or improved performance across CIFAR and ImageNet experiments.

  • Problem

    Overparameterized DNNs may generalize poorly, and SAM’s generalization benefits come with roughly double the computational cost of base optimizers.

  • Method

    ESAM combines Stochastic Weight Perturbation and Sharpness-sensitive Data Selection to approximate SAM’s sharpness using stochastic weights and selected sharpness-sensitive samples.

  • Results

    ESAM improves efficiency over SAM while preserving or improving test accuracy across CIFAR10, CIFAR100, and ImageNet experiments.

  • Takeaways & Limitations

    ESAM provides a more computationally efficient alternative to SAM without sacrificing its reported generalization performance.

  • Takeaways & Limitations

    SWP’s computational savings are reduced by chain-rule gradients for parameters outside the perturbed subset, especially in deeper networks.

Abstract

from arXiv · show

Overparametrized Deep Neural Networks (DNNs) often achieve astounding performances, but may potentially result in severe generalization error. Recently, the relation between the sharpness of the loss landscape and the generalization error has been established by Foret et al. (2020), in which the Sharpness Aware Minimizer (SAM) was proposed to mitigate the degradation of the generalization. Unfortunately, SAM s computational cost is roughly double that of base optimizers, such as Stochastic Gradient Descent (SGD). This paper thus proposes Efficient Sharpness Aware Minimizer (ESAM), which boosts SAM s efficiency at no cost to its generalization performance. ESAM includes two novel and efficient training strategies-StochasticWeight Perturbation and Sharpness-Sensitive Data Selection. In the former, the sharpness measure is approximated by perturbing a stochastically chosen set of weights in each iteration; in the latter, the SAM loss is optimized using only a judiciously selected subset of data that is sensitive to the sharpness. We provide theoretical explanations as to why these strategies perform well. We also show, via extensive experiments on the CIFAR and ImageNet datasets, that ESAM enhances the efficiency over SAM from requiring 100% extra computations to 40% vis-a-vis base optimizers, while test accuracies are preserved or even improved.

1 INTRODUCTION

ESAM is proposed to reduce SAM’s computational burden while preserving or improving its generalization performance. It combines stochastic weight perturbation and sharpness-sensitive data selection, with experiments reporting improved efficiency across datasets and architectures.

  • Overparameterized DNNs can achieve strong performance but remain prone to poor generalization when parameters greatly exceed training samples.
  • SAM improves generalization by seeking flatter minima, but its extra forward and backward operations substantially increase computational cost.
  • ESAM combines Stochastic Weight Perturbation and Sharpness-sensitive Data Selection to reduce SAM’s overhead while preserving comparable performance.SWP estimates sharpness from a stochastic weight neighborhood, while SDS selects samples whose losses are most sensitive to perturbation.
  • Experiments evaluate ESAM on CIFAR10, CIFAR100, and ImageNet across five DNN architectures, reporting 40% extra overhead versus base optimizers instead of SAM’s 100%.
  • ESAM integrates both strategies and is designed to improve DNN generalization with only marginal additional computational cost over standard training.

2 METHODOLOGY

SAM improves generalization by seeking flat minima, but its two-pass optimization doubles computation. ESAM reduces this overhead through stochastic weight perturbation and sharpness-sensitive data selection while preserving the intended SAM behavior.

  • Sharpness-Aware Minimization: SAM minimizes the maximum empirical-loss increase within a radius-ρ perturbation, encouraging solutions with lower sharpness and flatter minima.The sharpness is the maximal change in empirical loss after a norm-constrained weight perturbation.
  • Sharpness-Aware Minimization: Two forward and two backward operations make SAM require 100% extra computational overhead relative to base optimizers.The additional operations estimate the perturbation before the base optimizer updates the weights.
  • Stochastic Weight Perturbation: SWP randomly selects weights for perturbation, scaling the resulting partial gradient so its expected perturbation matches SAM’s.The selected subset is controlled by probability β, reducing the backward-propagation cost of sharpness estimation.
  • Stochastic Weight Perturbation: SWP’s computational savings are correlated with 1 −β, but chain-rule dependencies retain some backward cost, especially in deeper networks.The perturbation-selection probability β must therefore be tuned against generalization performance comparable to SAM’s.
  • Sharpness-Sensitive Data Selection: SDS selects samples whose losses increase most under the perturbation, forming a sharpness-sensitive subset B+ controlled by threshold α or ratio γ.The selected samples are used to approximate the perturbed full-batch objective with fewer samples.
  • Sharpness-Sensitive Data Selection: SDS uses LB+(fθ+ˆϵ) as a proxy for LB(fθ+ˆϵ), because B+ better matches the full batch’s gradients and improves efficiency without performance degradation.Experiments report higher gradient cosine similarity for B+ than for B−.

3 EXPERIMENTS

Experiments across CIFAR-10, CIFAR-100, and ImageNet show that ESAM reduces SAM’s training overhead while preserving or improving classification accuracy. Ablations indicate that SWP and SDS each contribute efficiency gains, and loss-landscape visualizations show flatter minima for SAM and ESAM than SGD.

  • CIFAR-10 and CIFAR-100: ESAM increases training speed by up to 40.30% compared with SAM across CIFAR-10 and CIFAR-100 experiments.The strongest efficiency result is 140.3% of SAM’s training speed for CIFAR-10 with ResNet-18.
  • CIFAR-10 and CIFAR-100: 85.56% accuracy is achieved by ESAM versus 84.46% for SAM on CIFAR-100 with PyramidNet-110.
  • ImageNet: ESAM outperforms SAM by 0.35% to 0.49% in accuracy and trains 28.7% faster on ImageNet.The ImageNet results are reported for ResNet-50 and ResNet-101.
  • Overall results: Across benchmark datasets and widely used DNN architectures, ESAM improves training speed and classification accuracy compared with SAM.
  • Ablation and parameter studies: SWP improves SAM’s training speed by 8.3% to 10.5% while achieving better performance, whereas SDS improves efficiency by 21.5% to 25.8%.SWP performs best with β = 0.6 or 0.5 depending on architecture, while SDS uses γ = 0.5 as the smallest value maintaining comparable performance.
  • Loss-landscape visualization: Both SAM and ESAM produce significantly flatter loss landscapes than SGD on ImageNet.The landscapes average 100 × 100 sampled points across ten groups of random Gaussian perturbations.
  • Ablation and parameter studies: SWP and SDS reduce computational overhead while achieving comparable or better performance than SAM, with β and γ allowing efficiency-performance trade-offs.

4 RELATED WORK

Research on sharpness connects flat minima with improved generalization, motivating SAM and subsequent extensions across training strategies, architectures, and tasks. The paper situates its loss-landscape analysis alongside this broader line of work.

  • Loss-landscape visualizations compare cross-entropy landscapes for ResNet50 trained with SGD, SAM, and ESAM on ImageNet.
  • Flat minima have been linked to better generalization and reduced overfitting through theoretical and empirical studies.
  • SAM quantifies sharpness through a maximization problem and uses it to train DNNs with improved generalization.
  • Subsequent work extends SAM-related sharpness regularization to adversarial training, other strategies, architectures, and tasks.

5 CONCLUSION

ESAM improves vanilla SAM’s efficiency by combining stochastic weight perturbation with sharpness-sensitive data selection, while preserving its generalization-oriented objective. The paper evaluates these strategies across datasets and architectures and identifies further overhead reduction as future work.

  • ESAM integrates Stochastic Weight Perturbation and Sharpness-sensitive Data Selection to improve vanilla SAM’s efficiency.
  • Both ESAM and SAM use two training steps: sharpness estimation followed by weight updating.
  • The strategies are theoretically motivated and evaluated across multiple datasets and DNN architectures.
  • Future work targets reducing ESAM’s computational overhead to the level of base optimizers by combining its two steps.

A.1 THE ALGORITHM OF SAM

The SAM algorithm contrasts with SGD by first computing a weight perturbation from the batch gradient, then evaluating a perturbed gradient for the weight update. This extra perturbation computation creates additional forward and backward work.

  • SAM samples a minibatch from the training set before computing the update.
  • SGD uses zero perturbation, whereas SAM computes a perturbation from the batch loss gradient.
  • SAM evaluates the gradient at perturbed weights and updates the weights using that gradient.
  • SAM’s perturbation calculation adds an extra forward-backward computation relative to SGD.

A.2 OPTIMIZING OVER SUBSET B+ IS REPRESENTATIVE

Sharpness-sensitive selection constructs B+ from examples whose loss changes most under the perturbation, using gradient-based reasoning to justify its representativeness. Experiments report that B+ better matches full-batch gradients than alternatives.

  • B+ is constructed by sorting each example’s loss change under the weight perturbation.
  • A first-order Taylor approximation expresses the loss change as the perturbation’s inner product with the per-example gradient.
  • The perturbation aggregates gradients across the complete dataset batch, linking example loss changes to gradient information.
  • Experiments verify the theoretical loss relationships across every training epoch.
  • B+ has higher cosine similarity with full-batch gradients than B− and a random subset.

A.3 LINEARITY MEASUREMENT OF SWP

SWP improves the linearity of the loss function used for sharpness approximation, with smaller β producing better linearity while β = {0.5, 0.6} balances accuracy and efficiency.

  • Accuracy–efficiency trade-off: SWP can improve ESAM accuracy compared with SAM, according to the reported experiments.
  • Linearity measurement: SWP improves the linearity of the loss function used to approximate the inner maximization in SAM.The linearity measure is ζ(ϵ, B) = |LB(fθ+ϵ) − LB(fθ) − ϵ⊤∇θLB(fθ)|.
  • Linearity measurement: Smaller β produces better linearity for SWP on ResNet18 trained on CIFAR10.
  • Accuracy–efficiency trade-off: β = {0.5, 0.6} provides the observed balance between ESAM accuracy and efficiency.Decreasing β improves linearity but reduces the magnitude of ˆϵ, worsening the inner maximization.

A.4 REDUCED COMPUTATIONAL OVERHEAD CONTRIBUTED BY SWP

SWP reduces computational overhead by perturbing only a stochastic subset of weights, with savings determined by selection probability, network depth, and architecture connectivity.

  • SWP overhead formulation: SWP selects basic parameter units for gradient computation with probability β, reducing the number of perturbed parameters and associated overhead.The saved overhead includes gradient calculation, storing, and hooking operations.
  • SWP overhead formulation: For architectures with multiple connected parameter units, the calculation-free rate compounds according to each layer’s number of parameter units.The formulation accounts for architecture-specific connectivity through K(n) and pn.
  • SWP overhead formulation: For serial architectures with one basic parameter unit per layer, the calculation-free rate follows pn = (1 − β)^n.The saved calculation from each parameter unit is summed across the network layers.
  • Toy validation: R2 = 0.9983 and R2 = 0.9989 validate the reduced-overhead formulations in the special and general MLP cases, respectively.The toy experiments used CIFAR10 MLPs with different layer connectivity patterns.

A.5 VISUALIZATION OF LOSS LANDSCAPES WITH RESPECT TO ADVERSARIAL WEIGHT PERTURBATIONS

Loss landscapes visualize sharpness under two orthogonal adversarial weight perturbations, showing that SAM and ESAM produce flatter minima than SGD on CIFAR10.

  • Landscape construction: The loss landscapes compare SGD, SAM, and ESAM under two orthogonal adversarial weight perturbations on CIFAR10.The axes represent η∇θLBx(fθ) and η∇θLBy(fθ), with Bx and By forming complementary random subsets of batch B.
  • Landscape comparison: SAM and ESAM improve sharpness significantly compared with SGD in the visualized CIFAR10 landscapes.
  • Architecture scope: ESAM was also evaluated with ViT-S/16 on ImageNet to assess efficiency in a Vision Transformer architecture.The reported comparison uses classification accuracy and training speed, with β = 0.5 and γ = 0.7.
  • Experimental settings: Training comparisons used tuned optimizer parameters, while ImageNet settings were adapted under computational-resource constraints.
Loading 2110.03141v2…