Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SAMPa: Sharpness-aware Minimization Parallelized

Wanyun Xie, Thomas Pethick, Volkan Cevher

TL;DR

A simple modification of SAM, termed SAMPa, is proposed, which allows to fully parallelize the two gradient computations, and achieves a twofold speedup of SAM under the assumption that communication costs between devices are negligible.

Abstract

Sharpness-aware minimization (SAM) has been shown to improve the generalization of neural networks. However, each SAM update requires \emph{sequentially} computing two gradients, effectively doubling the per-iteration cost compared to base optimizers like SGD. We propose a simple modification of SAM, termed SAMPa, which allows us to fully parallelize the two gradient computations. SAMPa achieves a twofold speedup of SAM under the assumption that communication costs between devices are negligible. Empirical results show that SAMPa ranks among the most efficient variants of SAM in terms of computational time. Additionally, our method consistently outperforms SAM across both vision and language tasks. Notably, SAMPa theoretically maintains convergence guarantees even for \emph{fixed} perturbation sizes, which is established through a novel Lyapunov function. We in fact arrive at SAMPa by treating this convergence guarantee as a hard requirement -- an approach we believe is promising for developing SAM-based methods in general. Our code is available at \url{https://github.com/LIONS-EPFL/SAMPa}.

SAMPa: Sharpness-aware Minimization Parallelized

TL;DR

A simple modification of SAM, termed SAMPa, is proposed, which allows to fully parallelize the two gradient computations, and achieves a twofold speedup of SAM under the assumption that communication costs between devices are negligible.

Abstract

Sharpness-aware minimization (SAM) has been shown to improve the generalization of neural networks. However, each SAM update requires \emph{sequentially} computing two gradients, effectively doubling the per-iteration cost compared to base optimizers like SGD. We propose a simple modification of SAM, termed SAMPa, which allows us to fully parallelize the two gradient computations. SAMPa achieves a twofold speedup of SAM under the assumption that communication costs between devices are negligible. Empirical results show that SAMPa ranks among the most efficient variants of SAM in terms of computational time. Additionally, our method consistently outperforms SAM across both vision and language tasks. Notably, SAMPa theoretically maintains convergence guarantees even for \emph{fixed} perturbation sizes, which is established through a novel Lyapunov function. We in fact arrive at SAMPa by treating this convergence guarantee as a hard requirement -- an approach we believe is promising for developing SAM-based methods in general. Our code is available at \url{https://github.com/LIONS-EPFL/SAMPa}.

Paper Structure

This paper contains 33 sections, 29 equations, 4 figures, 14 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background and Challenge of SAM
  3. SAM and its challenge
  4. Naive attempts
  5. Toy example.
  6. SAM Parallelized (SAMPa)
  7. SAMPa with stochasticity.
  8. Analysis
  9. Proofs for \ref{['sec:analysis']}
  10. Experiments
  11. Image classification
  12. CIFAR-10/100.
  13. ImageNet-1K.
  14. Efficiency comparison with efficient SAM variants
  15. Transfer learning
  16. ...and 18 more sections

Figures (4)

  • Figure 1: Comparison on $f(x)=\|x\|^2$.
  • Figure 2: Computational time comparison for efficient SAM variants. SAMPa-0.2 requires near-minimal computational time in both ideal and practical scenarios.
  • Figure 3: Test accuracy curve obtained from SAMPa algorithm using a range of $\lambda$.
  • Figure 4: Difference between $\nabla f(x_t)$ and $\nabla f(y_t)$.

Theorems & Definitions (1)

  • Remark 3