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.

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

Parallelizing Sharpness-Aware Minimization: A Semi-asynchronous, Small-Batch Approach

  • Weimin Li,
  • Yuzhong Sun

摘要

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.