Source-linked AI summary
ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, In Kwon Choi
TL;DR
Sharpness-based measures can be sensitive to parameter re-scaling that leaves the loss unchanged, weakening their connection with the generalization gap. The paper introduces scale-invariant adaptive sharpness and ASAM, which adjusts maximization regions accordingly; experiments report improved generalization performance across benchmark tasks. Adaptive sharpness also shows stronger correlation with the generalization gap than sharpness does.
Problem
Sharpness measured in a fixed-radius region is sensitive to parameter re-scaling that does not change the loss, weakening its connection with the generalization gap.
Method
The paper defines adaptive sharpness using normalization operators and proposes ASAM, which adaptively adjusts maximization regions with respect to weight scale.
Results
ASAM consistently improves generalization performance across image classification and machine translation tasks using various neural network architectures.
Takeaways & Limitations
Adaptive sharpness provides a scale-invariant generalization measure, while ASAM uses the corresponding bound to avoid the scale-dependency issue affecting SAM.
Takeaways & Limitations
The paper covers only element-wise and filter-wise normalization operators; layer-wise normalization is excluded because its required condition is not satisfied.
Abstract
from arXiv · showhide
Recently, learning algorithms motivated from sharpness of loss surface as an effective measure of generalization gap have shown state-of-the-art performances. Nevertheless, sharpness defined in a rigid region with a fixed radius, has a drawback in sensitivity to parameter re-scaling which leaves the loss unaffected, leading to weakening of the connection between sharpness and generalization gap. In this paper, we introduce the concept of adaptive sharpness which is scale-invariant and propose the corresponding generalization bound. We suggest a novel learning method, adaptive sharpness-aware minimization (ASAM), utilizing the proposed generalization bound. Experimental results in various benchmark datasets show that ASAM contributes to significant improvement of model generalization performance.
1. Introduction
The paper identifies scale-dependency as a weakness of sharpness-based generalization measures and introduces adaptive sharpness and ASAM to address it. Adaptive sharpness is scale-invariant, while ASAM uses it to adjust maximization regions and improve generalization across tasks.
- Motivation: Sharpness-based methods can assign different sharpness values to re-scaled parameters even when the loss and generalization gap remain unchanged.This scale-dependency may weaken the correlation between sharpness and generalization gap.
- Motivation: Prior scale-invariant generalization measures did not sufficiently investigate combining those measures with a learning algorithm.
- Contributions: Adaptive sharpness defines its maximization region using a normalization operator and remains unchanged under parameter re-scaling.The paper considers element-wise and filter-wise normalization operators.
- Contributions: ASAM adaptively adjusts maximization regions with respect to weight scale and minimizes a corresponding generalization bound.Its stated purpose is to avoid the scale-dependency issue affecting SAM.
- Contributions: Adaptive sharpness shows stronger correlation with generalization gap than ordinary sharpness in rank statistics.
- Experiments: ASAM consistently improves generalization performance across image classification and machine translation tasks using various neural network architectures.
2. Preliminary
The preliminary section defines the training and expected losses, introduces SAM as minimization of a PAC-Bayesian generalization bound, and describes its sharpness-based two-step optimization. SAM estimates a high-loss point in a fixed-radius region before updating the parameters using that point's gradient.
- Definitions: The generalization gap is the difference between expected loss on the data distribution and training loss, representing performance on unseen data.
- Sharpness-Aware Minimization: SAM aims to minimize a PAC-Bayesian generalization error upper bound based on loss-surface sharpness.
- Sharpness-Aware Minimization: SAM defines its maximization region as an ℓp ball with radius ρ for p ≥1, with sharpness determined over that fixed-radius region.
- Sharpness-Aware Minimization: Because the bound uses a strictly increasing function h, it can be replaced by an ℓ2 weight-decay regularizer with coefficient λ.
- Optimization Procedure: The two-step procedure follows from a first-order approximation and dual-norm formulation, with experiments reporting best performance at p = 2.
- Optimization Procedure: SAM iteratively estimates a loss-maximizing point near the current parameters, then performs descent at the current parameters using the gradient at that point.
3. Adaptive Sharpness: Scale-Invariant Measure of Generalization Gap
Adaptive sharpness addresses the scale-dependency of fixed-radius sharpness by adapting the maximization region through normalization operators. It is scale-invariant and shows stronger correlation with the generalization gap, while the paper evaluates element-wise and filter-wise variants and reports p = 2 as more suitable for learning.
- Scale-dependency problem: Fixed-radius sharpness can assign arbitrarily different values to re-scaled networks with identical generalization gaps, weakening their correlation.This scale-dependency arises because parameter re-scaling can leave the loss function unchanged while altering sharpness.
- Adaptive sharpness: Adaptive sharpness defines the maximization region using a normalization operator designed to cancel loss-preserving parameter scaling.The paper introduces normalization operators and defines adaptive sharpness from the resulting scale-adjusted region.
- Scale invariance: Theorem 1 establishes that adaptive sharpness has the same value at w and Aw for any invertible scaling operator A that leaves the loss unchanged.The corresponding maximization regions adapt under re-scaling, preserving the maximum loss in the illustrated examples.
- Empirical correlation: Adaptive sharpness shows stronger rank correlation with the generalization gap than sharpness across the reported hyper-parameter settings.The comparison uses mini-batch size, initial learning rate, weight decay coefficient, and dropout rate; the paper attributes the stronger correlation to scale invariance.
- Scope and choices: The paper covers element-wise and filter-wise normalization, excludes layer-wise normalization, and reports p = 2 as more suitable for learning than p = ∞.Layer-wise normalization does not satisfy the required condition, while node-wise alternatives are excluded because of assumptions or additional computational cost.
4. Adaptive Sharpness-Aware Minimization
ASAM formulates adaptive sharpness minimization as a minimax problem and approximates its inner maximization with a first-order update. The method uses a normalization operator and can be implemented with different optimization updates.
- Adaptive sharpness provides a generalization bound that remains valid beyond the standard p-norm setting, including p > 2.When the normalization operator is the identity, the bound reduces to the corresponding sharpness formulation.
- ASAM solves a minimax problem by first finding a perturbation that approximately maximizes the loss within the adaptive maximization region.The perturbation is obtained using a first-order approximation, analogously to SAM.
- The ASAM training procedure uses a two-step update that incorporates the adaptive perturbation before updating the model weights.The supplied algorithm description specifies iterative training over mini-batches, with radius, weight decay, and learning-rate parameters.
- Experiments evaluate ASAM with p = ∞ and p = 2, while the detailed algorithm uses p = 2.The SGD update can be combined with momentum or replaced by another optimization scheme such as Adam.
5. Experimental Results
Experiments evaluate ASAM on toy loss landscapes, image classification, machine translation, and label-noise settings. Across these tests, ASAM’s adaptive regions support scale-robust behavior and generally improve reported generalization outcomes.
- Toy Example: ASAM remains stable along equivalent loss contours, whereas SAM can fail to converge when the parameter scales and fixed radius are mismatched.With w0 = (0.3, 0.033), SAM fails to converge for ρ = 0.05, while ASAM converges for the tested initializations when ρ is sufficiently small.
- Normalization Factors: Element-wise normalization and filter-wise normalization produce comparable test accuracies across ρ, with element-wise normalization slightly better at ρ = 1.0.
- Normalization Factors: p = 2 yields better test accuracies than p = ∞, while bias normalization does not improve accuracy; subsequent tests therefore use element-wise normalization with p = 2 without bias normalization.
- Image Classification: ASAM achieves higher maximum test accuracies than SGD and SAM across the evaluated rectifier architectures and image-classification datasets beyond CIFAR.
- Machine Translation: Adam+ASAM improves BLEU on IWSLT’14 DE-EN compared with Adam and Adam+SAM using a Transformer architecture.
- Robustness to Label Noise: ASAM generally improves test accuracy over SGD and SAM across symmetric label-noise levels from 20% to 80%.
6. Conclusions
The conclusion presents adaptive sharpness as both a scale-invariant training principle and a generalization measure. ASAM improves generalization in the reported comparisons, while broader normalization schemes remain for future investigation.
- Adaptive sharpness adjusts maximization regions to weight scale, improving the training path while preserving scale invariance.
- ASAM shows superior performance to SAM in the reported image-classification comparison tests and improves generalization performance.
- Adaptive sharpness has stronger correlation with generalization gap and may serve as a neural-network assessment metric.
- The paper does not cover all normalization schemes satisfying its stated condition, leaving their investigation for future work.
A.1. Proof of Theorem 2
This appendix develops a concentration-inequality-based proof for a generalization bound under Gaussian perturbations. The proof uses an assumption that perturbation does not reduce expected loss and then bounds the resulting expression.
- The proof begins with a concentration inequality for independent normal variables with common variance σ^2.
- The maximum standard deviation is defined as σmax = max{σi}.
- The argument invokes the Laurent–Massart concentration result for any x > 0.
- The theorem assumes that adding Gaussian perturbation to normalized weights does not improve expected loss.
- Under that assumption, an existing perturbation-based generalization bound is applied and the left-hand side is subsequently bounded.
B. Correlation Analysis
The correlation analysis compares sharpness measures with generalization gap using rank-based statistics across a full hyper-parameter grid. It evaluates both aggregate and per-hyper-parameter correlations.
- Rank Correlation: Kendall rank correlation coefficient measures agreement between the rankings of a generalization measure and generalization gap, ranging from −1 to 1.
- Granulated Correlation: Granulated coefficients isolate correlation with respect to each hyper-parameter rather than aggregating all hyper-parameter effects.
- Experimental Design: The study varies mini-batch size, initial learning rate, weight decay coefficient, and dropout rate, using five values for each and producing 625 configurations.
- Experimental Design: The configurations are used to train WideResNet-28-2 on CIFAR-10 with SGD for 200 epochs and cosine learning-rate decay.
- Measure Comparison: Adaptive sharpness is compared with sharpness under element-wise normalization, using p = 2 and p = ∞ and grid-searched ρ values selected for maximum correlation.