Asynchronous Sharpness-Aware Minimization For Fast and Accurate Deep Learning

arXiv:2503.11147 · cs.LG · Submitted 2025-03-14 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Asynchronous Sharpness-Aware Minimization For Fast and Accurate Deep Learning".

Jane: The gist The proposed method enables to fully utilize heterogeneous system resources such as CPUs and GPUs by designing a general optimization algorithm that improves the model’s generalization performance virtually…

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So wrapping up this discussion on "Asynchronous Sharpness-Aware Minimization For Fast and Accurate Deep Learning," the authors are showing that they built an asynchronous-parallel SAM algorithm that manages to keep the gradient norm penalizing effect nearly identical to the original while completely hiding the model perturbation time by adjusting batch size system-aware.

Jane: The title suggests a focus on speed and accuracy, and what we see is a method that achieves those goals by leveraging asynchronous computation and hardware specifics, which is really about making SAM practical for real-world applications.

Lu: The implication is that we can use heterogeneous system resources like CPUs and GPUs much more effectively during neural network training without sacrificing the generalization performance benefits of SAM.

Meng: For an engineer looking at this, it means a potential way to reduce the compute budget for high-quality model fine-tuning tasks by optimizing how we run those gradient steps on our actual machines.

Lalam: Essentially, they're providing a practical path to use these complex optimization techniques in real training environments instead of just theoretical settings.

Tom: And the paper concludes that this asynchronous SAM is a practical method that is readily applicable to real-world applications, efficiently improving their generalization performance, as stated on page two of THIS PAPER <ref:2503.11147#pg2,asynchronous SAM is a practical method that is readily applicable to real>.

Conclusion: Tom: So we’ve been talking about how these authors took Sharpness-Aware Minimization, or SAM—which is this optimization trick to make models generalize better—and made it way faster using some clever scheduling tricks.

Jane: Right, it’s about taking a method that’s theoretically solid and making it practical for actual training without needing huge amounts of time. The title itself says "Asynchronous" because they're changing how the model learns step by step.

Lu: What they did is design a general optimization algorithm that uses heterogeneous resources, like mixing CPUs and GPUs, to fully use them instead of letting one resource bottleneck everything.

Meng: So the big picture here is that they’ve figured out a way to improve performance without necessarily increasing the wall-clock training time for the model itself. That’s a huge deal for deployment speed.

Lalam: From my view, this means we can train bigger, more complex AI models on less hardware because we can hide the slow parts of the learning process behind other fast computations.

Tom: Exactly. The authors show that they get almost the same generalization accuracy as the original SAM method, even when using different hardware setups. It’s not just about being faster; it’s about keeping that high performance while running on real machines.

Jane: And they do this by adjusting the batch size in a way that depends on whether you're using a slow CPU or a fast GPU to balance the time spent calculating gradients versus updating the model weights.

Lu: It’s really interesting how they define this asynchrony term, tau, which basically lets them deliberately use slightly older information for one part of the process while updating with newer stuff in another.

Meng: From an engineering standpoint, that system-aware batch size adjustment is what makes it work on real-world setups like having a few CPUs and some GPUs together. It automatically finds the best balance between calculation time and update time.

Lalam: I think the implication for culture is that we can start building models that are smarter without needing an infinitely powerful supercomputer just to run the optimization steps themselves.

Tom: It sounds like these authors have really nailed this trade-off between speed and accuracy in a way that’s directly applicable right now. But there’s still one thing they pointed out.

Jane: Yeah, they did mention a limitation, so remember, it says that if you use the batch size too small for this asynchronous approach, you can lose some of that accuracy gain.

Lu: So it’s not a magic fix for everything; there’s still a tuning part to make sure those parameters are right. It shows the complexity of making these things truly efficient across different systems.

Meng: So, if you're building a training pipeline, you can look at this as a way to squeeze more performance out of your existing hardware configuration without needing massive upgrades just for the optimization process.

Inha University

cs.LG

Submitted: 2025-03-14

Updated: 2026-10-13

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 91/100

The gist: The gist The proposed method enables to fully utilize heterogeneous system resources such as CPUs and GPUs by designing a general optimization algorithm that improves the model’s generalization

Key concepts

Sharpness-Aware Minimization (SAM): SAM is an optimization technique that improves a model's generalization by calculating gradients twice per iteration: once for updating the model and once for perturbing it. This perturbation helps find flatter minima, leading to better performance on unseen data.
Asynchronous Model Perturbation: This involves deliberately introducing a slight delay or asynchrony when calculating the gradient used to perturb the model. Instead of using the absolute latest parameters, it uses slightly older ones. This breaks the direct dependency between perturbation and update steps, making it computationally cheaper.
System-Aware Batch Size Adjustment: This technique dynamically adjusts the batch size based on hardware capabilities. It calculates a new batch size by scaling the original by the ratio of fast compute time to slow compute time. This ensures that the time spent calculating gradients for both perturbation and update is nearly equal, maximizing resource utilization.

Terminology

Generalization — 2× in this paper · explained in 13 episodes
Sharpness-Aware Minimization — 2× in this paper
Accuracy — 1× in this paper · explained in 1 episode
Fine-Tuning — 1× in this paper · explained in 5 episodes
Transformer — 1× in this paper · explained in 2 episodes

Summary

The gist The proposed method enables to fully utilize heterogeneous system resources such as CPUs and GPUs by designing a general optimization algorithm that improves the model’s generalization performance virtually not increasing the training time

Key Motivation and Problem

Sharpness-Aware Minimization (SAM) is an optimization method designed to improve machine learning models’ generalization performance Specifically, it calculates gradients twice for a single training iteration – one for model perturbation (gradient ascent) and another for model update (gradient descent) To make SAM a more compelling option for deep learning applications, thus, addressing this high computational cost is crucial

Asynchronous Model Perturbation

The paper studies how to take advantage of such an asynchronous scheme to break the data dependency between the model perturbation and the model update steps in SAM and thus tackle the expensive model perturbation cost issue Specifically, we propose deliberately injecting a limited degree of asynchrony into the gradient ascent to perturb and update the model simultaneously Instead of calculating the gradient for the model perturbation using the most recent model parameters, our approach uses slightly out-of-date parameters This asynchronous approach makes a practical trade-off between the convergence rate and the system utilization The model perturbation can be defined as follows: wt+1 = wt − η∇Lt−τ (wt−τ)∥∇Lt−τ (wt−τ)∥ The key difference between (2) and the original SAM is the degree of asynchrony term τ on the right-hand side If τ > 0, it implies that the model is perturbed using the slightly out-of-date gradients 0, it implies that the model is perturbed using the slightly out-of-date gradients>

System-Aware Batch Size Adjustment

The study investigates the possibility of utilizing these CPU resources to perturb the model in SAM by reducing the gradient ascent batch size This enables our asynchronous-parallel SAM to entirely hide the model perturbation time behind the model update time In modern computers, it is common to have one or two sockets of CPUs together with a few accelerators like GPUs We propose to adjust b′ in a system-aware manner such that b′ = (Tf /Ts) × b Here, Ts is the single data gradient calculation time on the slow compute resource and Tf is that on the fast compute resource In this way, we can automatically find the maximum batch size b′ which makes the gradient ascent time and the gradient descent time nearly the same

Performance Results

Extensive experiments demonstrate that regardless of the heterogeneous system configurations, asynchronous SAM eliminates the model perturbation time during training while maintaining its superior generalization performance For instance, asynchronous SAM achieves 92.60% CIFAR-10 accuracy that is almost the same as the original SAM’s accuracy, 92.53% Our proposed method achieves clearly higher accuracy than SGD in all benchmarks Furthermore, asynchronous SAM shows almost the same accuracy as that of Generalized SAM which achieves the best validation accuracy among all the SOTA methods The results empirically prove that asynchronous SAM enables users to exploit the in-node heterogeneous system resources to improve the generalization performance

Convergence Analysis

The convergence rate of Algorithm 1 is analyzed as follows This novel approach enables users to exploit heterogeneous system resources such as CPU and GPU to accelerate the neural network training with SAM Our extensive experiments well demonstrate that the proposed method effectively reduces the training time while maintaining the superior generalization performance of SAM

Limitations and Future Work

Our experimental results show that the asynchronous SAM loses the accuracy when the gradient ascent batch size is too small<ref:2503.

Improvements for AI systems

Asynchronous Parallel SAM for Heterogeneous Systems: Implement a training pipeline where gradient ascent runs on a slow compute resource (e.g., CPU) using a smaller batch size, while gradient descent runs on a fast resource (e.g., GPU) with the full batch size, as described in Section 3.3 and Algorithm 1, to entirely hide the model perturbation time behind the model update time.
System-Aware Batch Size Adjustment: Automatically adjust the gradient ascent batch size, denoted as b', using a system-aware formula like b' = (Tf / Ts) × b, where Tf and Ts are the data gradient calculation times on fast and slow resources, to ensure the gradient ascent time and the gradient descent time nearly the same.
Gradient Stability Exploitation: Leverage the observation that gradients are stable during training (The similarity remains high during the training regardless of which data and model are used) by using slightly out-of-date parameters for perturbation, as defined in Equation (2) with a fixed staleness term τ = 1, to achieve a model perturbation and the model update simultaneously.
Resource Utilization Synergy: Harmonize the proposed asynchronous SAM with gradient recycling approaches (Asynchronous SAM is orthogonal to the gradient recycling approach) to better exploit extremely heterogeneous system resources such as CPU and GPU in modern computing environments.

Abstract

Sharpness-Aware Minimization (SAM) is an optimization method that improves the generalization performance of machine learning models. Despite its effectiveness, it has not been widely adopted in real-world applications due to its high computational cost. Existing approaches reduce this overhead by approximating or reusing gradients in the perturbation step, but still assume sequential perturbation and update steps. In this study, we introduce AsyncSAM, a variant of SAM that enables concurrent model perturbation and parameter updates across multiple compute resources. By adjusting the perturbation batch size in a system-aware manner, our asynchronous scheme fully overlaps perturbation with updates, eliminating the overhead and allowing even CPUs to perform perturbations without extra time. We also provide theoretical analysis showing that AsyncSAM achieves the same order of convergence as conventional SAM. Extensive experiments demonstrate that AsyncSAM delivers superior generalization performance and the highest training throughput among state-of-the-art SAM variants. For example, in ViT fine-tuning on CIFAR-100, it achieves accuracy comparable to SAM while maintaining training time nearly identical to SGD.