Asymptotic Unbiased Sample Sampling to Speed Up Sharpness-Aware Minimization

TL;DR

This paper introduces AUSAM, a probabilistic sampling method that accelerates Sharpness-Aware Minimization by reducing computational cost while maintaining generalization performance across various tasks and models.

Contribution

AUSAM is a novel, architecture-agnostic sampling technique that approximates gradient norms to speed up SAM without sacrificing accuracy.

Findings

AUSAM achieves over 70% speedup on CIFAR datasets.

AUSAM maintains comparable accuracy to SAM across tasks.

AUSAM outperforms recent data pruning methods in efficiency and performance.

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.