Parallelizing Sharpness-Aware Minimization: A Semi-asynchronous, Small-Batch Approach
摘要
Sharpness-aware minimization (SAM), which seeks flat minima of the loss landscape, has proven effective in improving the generalization performance of deep learning models. However, SAM inherently calculates two gradients sequentially per step, doubling the training overhead while leaving the potential to improve performance. In this paper, we propose a novel approach that enhances both training efficiency and the generalization performance of SAM. Our approach leverages the noteworthy observations that (1) Asynchronous SGD gradient with a small delay can align with SAM gradient in direction while retaining efficient convergence, yet it struggles to match its scale and achieve comparable generalization. (2) Adopting smaller batch sizes in ASGD bolsters generalization with little impact on convergence but decreases throughput. Based on these insights, we propose ParaS3AM, a parallel algorithm that approximates SAM using one-step delayed ASGD while enhancing generalization performance using smaller batch sizes without substantial losses in throughput. Experiments on various datasets and architectures demonstrate that ParaS3AM outperforms state-of-the-art baselines and achieves the optimal performance-efficiency trade-off.