Source-linked AI summary
EMASAM: a Computationally Efficient Sharpness-Aware Minimization via EMA-Guided Perturbations
Tanapat Ratchatorn, Masayuki Tanaka
TL;DR
SAM improves generalization but doubles backward propagation cost, limiting practicality. EMASAM replaces gradient-based perturbations with EMA-guided directions and, across architectures and benchmark datasets, consistently improves generalization and overall performance while eliminating extra backpropagation.
Problem
SAM and its variants improve generalization but require additional backward propagations, increasing computational overhead and limiting practicality for large-scale models and datasets.
Method
EMASAM uses the discrepancy between EMA shadow-model parameters and current parameters to define a stable, gradient-free perturbation direction.
Results
Across various neural network architectures and benchmark datasets, EMASAM consistently improves generalization and overall performance over standard SAM while eliminating an extra backpropagation.
Takeaways & Limitations
EMASAM provides a computationally efficient alternative to SAM that preserves its generalization and robustness while reducing perturbation-related gradient noise.
Takeaways & Limitations
SAM’s gradient-based perturbation direction can vary significantly across iterations, potentially destabilizing the perturbation step and training dynamics.
Abstract
from arXiv · showhide
Recent progress in optimization research has highlighted the sharpness of the loss landscape as a key factor in narrowing the generalization gap. Motivated by this insight, Sharpness-Aware Minimization (SAM) was proposed as a training strategy that enhances generalization. Despite the promising performance, SAM suffers from its twice computational cost due to its core algorithm requiring an extra gradient computation during the perturbation step. To overcome this limitation, we introduce Exponential Moving Average Sharpness-Aware Minimization (EMASAM), a computationally efficient variant of SAM. EMASAM does not require the loss gradient in the perturbation step. Instead, EMASAM defines the perturbation direction based on the discrepancy between the main model and the EMA shadow model. This perturbation travels away from the stable average position toward the less stable area, acting as a softer yet cheaper alternative to SAM's worst-case scenario perturbation. Moreover, since EMASAM's perturbation does not rely on noisy mini-batch gradients, it mitigates the gradient-induced instability inherent in SAM. Hence, EMASAM eliminates the need for an extra backpropagation while also preserving the generalization ability of the SAM-style training. Several experiments have been performed and confirm the efficiency and robustness of our method.
1 Introduction
The introduction frames generalization and loss-landscape sharpness as central concerns, while identifying SAM’s computational cost and gradient-noise sensitivity as limitations. It presents EMASAM, which uses EMA-based perturbations to remove SAM’s extra backward pass while preserving perturbation stability and generalization.
- Motivation: Overparameterized models can overfit, creating a performance gap between training and testing data and making generalization a critical objective.
- Motivation: Loss-surface sharpness and flatter minima are presented as important factors in model generalization, motivating optimization methods that seek flatter solutions.
- SAM limitations: SAM improves generalization but requires two backward propagations per training iteration, doubling training time and computational overhead compared with standard training.
- SAM limitations: SAM’s gradient-normalized perturbation direction is sensitive to mini-batch noise, which can make it unstable or inconsistent across training steps.
- EMASAM: EMASAM integrates historically smoothed EMA weights into the perturbation process instead of relying on SAM’s gradient-based perturbation.
- EMASAM: EMASAM eliminates SAM’s extra backward pass while preserving perturbation stability, generalization, and robustness, with experiments reporting improved generalization and overall performance over standard SAM.
2 Preliminary
Conventional training can converge to sharp minima that generalize poorly, whereas SAM seeks parameters with uniformly low loss in a neighborhood. EMA stabilizes training by averaging model parameters, providing a smoother solution without changing the objective or adding backward passes.
- Sharp minima can have low training loss but substantially higher loss on unseen data, leading to weaker generalization than flatter solutions.
- SAM minimizes perturbed loss over a parameter neighborhood, steering optimization toward flatter minima associated with improved generalization.Its perturbation is bounded by ρ in the L2-norm, and the method evaluates gradients at the perturbed position for weight updates.
- For small ρ, SAM’s perturbation follows the direction where the loss increases most sharply, while StopGrad keeps the perturbation fixed during weight-gradient computation.
- EMA maintains an exponentially weighted shadow model that filters stochastic-update noise and captures a more stable parameter representation.The shadow model is updated after each optimizer step, with decay rate α controlling the influence of past parameters.
- EMA can approximate a smoother, more generalizable point in the loss landscape without modifying the optimization objective or requiring additional backward passes.
3 Proposed Method
EMASAM is a computationally efficient SAM variant that replaces gradient-based perturbations with EMA-guided directions from a stable shadow model toward the current model. This removes the extra gradient computation and reduces gradient-induced perturbation instability while preserving SAM-style robustness and generalization.
- EMA-guided perturbation: EMASAM replaces SAM’s gradient-based perturbation with a direction defined by the discrepancy between the EMA shadow model and the current model.The EMA model provides a temporally smoothed, stable reference, while the current model is less stable.
- EMA-guided perturbation: The EMA-guided vector moves parameters from a stable position toward a less stable configuration, serving as a softer approximation of SAM’s worst-case direction.Unlike SAM’s maximally adversarial perturbation, EMASAM moves toward a less favorable configuration.
- Stability: Because EMASAM’s perturbation does not rely on noisy mini-batch gradients, its direction avoids batch-to-batch fluctuations and is more stable than instantaneous gradient-based perturbations.The direction is defined relative to the globally referenced, smoothed EMA position.
- Update procedure: EMASAM computes the weight-update gradient at the perturbed position, then updates the original parameters using that gradient.This preserves the two-step structure of computing at the perturbed position and updating from the original position.
- EMA maintenance: The EMA shadow model is maintained after each update as a temporally smoothed version of the training model using a decay coefficient.Its parameters are averaged over time after each model update.
- Efficiency and robustness: EMASAM eliminates additional gradient computation while providing a more stable perturbation direction without extra computational overhead.The method is presented as maintaining generalization and robustness despite the reduced computation.
4 Experiments
Experiments across multiple architectures and datasets show that EMASAM generally matches or outperforms SAM, EMA, and other variants while requiring only one backpropagation. Its advantage persists across CIFAR-scale and ImageNet-1K evaluations, with perturbation direction contributing importantly to robustness.
- ResNet-18: On ResNet-18, EMASAM achieves the highest test accuracies on CIFAR-10, CIFAR-100, and Fashion-MNIST, ranks second on EMNIST, and leads average accuracy.Standard SAM slightly exceeds EMASAM on EMNIST.
- Additional architectures: Across WideResNet-28-10 and PyramidNet-110, EMASAM leads average accuracy despite trailing SAM variants on some WideResNet datasets and ranking second on PyramidNet EMNIST.EMASAM is best on CIFAR-10 and Fashion-MNIST for both architectures, and also on CIFAR-100 for PyramidNet-110.
- Method comparison: EMASAM often matches or outperforms EMA, standard SAM, and other SAM variants, whereas random EMA perturbations significantly degrade robustness.These results support the importance of EMASAM’s defined perturbation direction.
- Efficiency: EMASAM achieves the best test accuracy with only a single backpropagation, demonstrating efficient training without sacrificing robustness.The EMA random-perturbation result is omitted from the comparison figure because it is out of range.
5 Conclusion
EMASAM is presented as a computationally efficient alternative to SAM that reduces overhead by deriving perturbations from the discrepancy between current parameters and an EMA shadow model rather than loss gradients.
- EMASAM reduces computational overhead by deriving perturbation directions from the discrepancy between current model parameters and their EMA-based shadow, without loss-gradient information.The perturbation pushes the model away from a temporally stable reference toward a less stable configuration.