Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Asymptotic Unbiased Sample Sampling to Speed Up Sharpness-Aware Minimization

Jiaxin Deng, Junbiao Pang, Baochang Zhang

TL;DR

This work tackles the high computational cost of Sharpness-Aware Minimization (SAM) by proposing AUSAM, a sampling-based acceleration that selects mini-batch examples with large gradient-norm signals. The method uses a Difference in Loss (DLP) proxy for per-sample gradient influence and maintains an Average DLP (ADLP) history to guide probabilistic sampling, achieving asymptotic unbiasedness as training progresses. The authors provide theoretical analysis under standard smoothness and variance assumptions and demonstrate empirical speedups of about 70% on CIFAR-10/100 and Tiny-ImageNet while preserving SAM-level generalization; AUSAM is also shown effective for human pose estimation and quantization-aware training. Overall, AUSAM offers a plug-and-play, architecture-agnostic approach to accelerate SAM without sacrificing performance, with meaningful practical impact for large-scale models and varied computer-vision tasks.

Abstract

Sharpness-Aware Minimization (SAM) has emerged as a promising approach for effectively reducing the generalization error. However, SAM incurs twice the computational cost compared to base optimizer (e.g., SGD). We propose Asymptotic Unbiased Sampling with respect to iterations to accelerate SAM (AUSAM), which maintains the model's generalization capacity while significantly enhancing computational efficiency. Concretely, we probabilistically sample a subset of data points beneficial for SAM optimization based on a theoretically guaranteed criterion, i.e., the Gradient Norm of each Sample (GNS). We further approximate the GNS by the difference in loss values before and after perturbation in SAM. As a plug-and-play, architecture-agnostic method, our approach consistently accelerates SAM across a range of tasks and networks, i.e., classification, human pose estimation and network quantization. On CIFAR10/100 and Tiny-ImageNet, AUSAM achieves results comparable to SAM while providing a speedup of over 70%. Compared to recent dynamic data pruning methods, AUSAM is better suited for SAM and excels in maintaining performance. Additionally, AUSAM accelerates optimization in human pose estimation and model quantization without sacrificing performance, demonstrating its broad practicality.

Asymptotic Unbiased Sample Sampling to Speed Up Sharpness-Aware Minimization

TL;DR

This work tackles the high computational cost of Sharpness-Aware Minimization (SAM) by proposing AUSAM, a sampling-based acceleration that selects mini-batch examples with large gradient-norm signals. The method uses a Difference in Loss (DLP) proxy for per-sample gradient influence and maintains an Average DLP (ADLP) history to guide probabilistic sampling, achieving asymptotic unbiasedness as training progresses. The authors provide theoretical analysis under standard smoothness and variance assumptions and demonstrate empirical speedups of about 70% on CIFAR-10/100 and Tiny-ImageNet while preserving SAM-level generalization; AUSAM is also shown effective for human pose estimation and quantization-aware training. Overall, AUSAM offers a plug-and-play, architecture-agnostic approach to accelerate SAM without sacrificing performance, with meaningful practical impact for large-scale models and varied computer-vision tasks.

Abstract

Sharpness-Aware Minimization (SAM) has emerged as a promising approach for effectively reducing the generalization error. However, SAM incurs twice the computational cost compared to base optimizer (e.g., SGD). We propose Asymptotic Unbiased Sampling with respect to iterations to accelerate SAM (AUSAM), which maintains the model's generalization capacity while significantly enhancing computational efficiency. Concretely, we probabilistically sample a subset of data points beneficial for SAM optimization based on a theoretically guaranteed criterion, i.e., the Gradient Norm of each Sample (GNS). We further approximate the GNS by the difference in loss values before and after perturbation in SAM. As a plug-and-play, architecture-agnostic method, our approach consistently accelerates SAM across a range of tasks and networks, i.e., classification, human pose estimation and network quantization. On CIFAR10/100 and Tiny-ImageNet, AUSAM achieves results comparable to SAM while providing a speedup of over 70%. Compared to recent dynamic data pruning methods, AUSAM is better suited for SAM and excels in maintaining performance. Additionally, AUSAM accelerates optimization in human pose estimation and model quantization without sacrificing performance, demonstrating its broad practicality.
Paper Structure (18 sections, 5 theorems, 11 equations, 6 tables, 1 algorithm)

This paper contains 18 sections, 5 theorems, 11 equations, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Sharpness-Aware Minimization
  4. Methods of Accelerating SAM
  5. Method
  6. Principle of Sampling
  7. Efficiently Per-sample gradient vs. Directly Approximate Gradient Norms
  8. Sampling Strategy
  9. Speed up ratio of SAM
  10. Theoretical Analysis
  11. Experimental Results
  12. Setup
  13. Parameter Studies
  14. Comparison to SOTAs
  15. Comparison to Dynamic Data Pruning
  16. ...and 3 more sections

Key Result

Theorem 1

Suppose we select $N$ samples from a mini-batch $\mathbf{B}$ of size $K$, denoting them as $\hat{\mathbf{B}}$ ($\hat{{x}_i} \in \hat{\mathbf{B}}$, $i=1,2,...,N$), and denote the remaining set of $M(M=K-N)$ samples as $\check{\mathbf{B}}$ ($\check{{x}_i} \in \check{\mathbf{B}}$, $i=1,2,...,M$). Let $

Theorems & Definitions (5)

  • Theorem 1
  • Lemma 1
  • Theorem 2
  • Theorem 3
  • Theorem 4