PaperlayerView exact revision on arXiv

Source-linked AI summary

Surrogate Gap Minimization Improves Sharpness-Aware Training

Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvornek, Sekhar Tatikonda, James Duncan, Ting Liu

arXiv:2203.08065v2cs.LGcs.AI

TL;DR

SAM’s perturbed-loss objective can assign low values to both sharp and flat minima, leaving sharpness selection unresolved. The paper introduces the surrogate gap and GSAM, which jointly target low perturbed loss and low sharpness; GSAM reports better generalization than SAM and AdamW with negligible overhead. Its optimization combines objectives whose gradients can conflict, and its ethical impact depends on the models trained with it.

  • Problem

    SAM minimizes a perturbed loss, but low perturbed loss can occur at both sharp and flat minima, so it does not always select flat solutions.

  • Method

    GSAM minimizes perturbed loss and surrogate gap using gradient decomposition and an orthogonal ascent step that leaves perturbed loss unchanged.

  • Results

    GSAM consistently improves generalization over SAM and vanilla training across architectures; on ViT-B/32, it improves top-1 accuracy by +3.2% over SAM and +5.4% over AdamW.

  • Takeaways & Limitations

    The surrogate gap provides an easy-to-compute sharpness measure, while GSAM offers improved generalization over SAM at negligible computation cost.

  • Takeaways & Limitations

    Combining perturbed loss and surrogate gap through a weighted sum can be sub-optimal because their gradients may conflict and gradient descent on the gap can increase perturbed loss.

Abstract

from arXiv · show

The recently proposed Sharpness-Aware Minimization (SAM) improves generalization by minimizing a \textit{perturbed loss} defined as the maximum loss within a neighborhood in the parameter space. However, we show that both sharp and flat minima can have a low perturbed loss, implying that SAM does not always prefer flat minima. Instead, we define a \textit{surrogate gap}, a measure equivalent to the dominant eigenvalue of Hessian at a local minimum when the radius of the neighborhood (to derive the perturbed loss) is small. The surrogate gap is easy to compute and feasible for direct minimization during training. Based on the above observations, we propose Surrogate \textbf{G}ap Guided \textbf{S}harpness-\textbf{A}ware \textbf{M}inimization (GSAM), a novel improvement over SAM with negligible computation overhead. Conceptually, GSAM consists of two steps: 1) a gradient descent like SAM to minimize the perturbed loss, and 2) an \textit{ascent} step in the \textit{orthogonal} direction (after gradient decomposition) to minimize the surrogate gap and yet not affect the perturbed loss. GSAM seeks a region with both small loss (by step 1) and low sharpness (by step 2), giving rise to a model with high generalization capabilities. Theoretically, we show the convergence of GSAM and provably better generalization than SAM. Empirically, GSAM consistently improves generalization (e.g., +3.2\% over SAM and +5.4\% over AdamW on ImageNet top-1 accuracy for ViT-B/32). Code is released at \url{ https://sites.google.com/view/gsam-iclr22/home}.

1 INTRODUCTION

SAM improves generalization by minimizing a neighborhood-based perturbed loss, but that objective can favor both sharp and flat minima. GSAM adds surrogate-gap minimization through an orthogonal ascent step, with theoretical and empirical support for improved generalization at negligible overhead.

  • Entropy-SGD avoids sharp valleys but requires nested Langevin iterations with substantially higher computation costs than vanilla training.
  • SAM minimizes the maximum loss within a parameter neighborhood rather than the loss at a single point.
  • Low perturbed loss can occur at both sharp and flat minima, so minimizing it alone is not always sharpness-aware.
  • GSAM jointly minimizes perturbed loss and surrogate gap, using an orthogonal ascent step that does not change the perturbed loss.
  • GSAM is reported to have negligible computation overhead, convergence guarantees, provably better generalization than SAM, and validation across ResNets, Vision Transformers, and MLP-Mixers.

2 PRELIMINARIES

The preliminaries define the original, perturbed, and surrogate-gap losses and explain SAM’s neighborhood optimization. Gradient decomposition supplies the components later used to distinguish perturbed-loss minimization from sharpness control.

  • Notations: The perturbed loss fp(w) is the worst loss within a radius-ρt ball around parameter w at training step t.
  • Notations: The surrogate gap h(w) equals fp(w) − f(w), while ηt controls the learning rate and α scales GSAM’s ascent step.
  • Gradient decomposition: The original gradient is decomposed into components parallel and orthogonal to the perturbed-loss gradient by projection.
  • Sharpness-Aware Minimization: SAM replaces the original-loss objective with approximately minimizing the loss at an adversarially perturbed weight wadv.
  • Sharpness-Aware Minimization: The perturbation maximization is approximated by a first-order Taylor expansion when ρ is small.

3 THE SURROGATE GAP MEASURES THE SHARPNESS AT A LOCAL MINIMUM

The paper argues that perturbed loss does not reliably measure sharpness and introduces the surrogate gap as an equivalent local-minimum measure. GSAM uses gradient decomposition to reduce this gap while preserving perturbed-loss optimization.

  • The perturbed loss is not always sharpness-aware: SAM is not guaranteed to find flat solutions, with sharpness measured here by the Hessian’s dominant eigenvalue σmax at a local minimum.
  • GSAM optimization: GSAM projects gradients into parallel and orthogonal components, then uses the orthogonal component for its ascent update while preserving the perturbed-loss direction.
  • The perturbed loss is not always sharpness-aware: A lower perturbed loss does not necessarily imply a lower dominant Hessian eigenvalue, so SAM can prefer a sharper minimum.
  • The surrogate gap agrees with sharpness: The surrogate gap h(w) is the difference between maximum neighborhood loss and center-point loss, and it is equivalent to local sharpness for sufficiently small ρ.
  • The surrogate gap agrees with sharpness: Because h is non-negative and decreases toward zero on flatter surfaces, minimizing it can identify regions with flatter loss surfaces.

4 SURROGATE GAP GUIDED SHARPNESS-AWARE MINIMIZATION

GSAM jointly minimizes SAM’s perturbed loss and a surrogate gap, using an orthogonal ascent update to reduce sharpness without changing the perturbed loss.

  • General idea: GSAM simultaneously minimizes the perturbed loss fp and surrogate gap h to seek low-loss regions with flatter surfaces.Minimizing fp targets low losses within the neighborhood, while minimizing h targets flat local minima.
  • General idea: The surrogate-gap objective can conflict with perturbed-loss minimization when directly combined as fp(w)+λh(w), potentially increasing fp.The paper attributes this conflict to opposing gradients ∇fp(w) and ∇h(w).
  • Gradient decomposition and ascent: GSAM decomposes ∇f(w) into components parallel and orthogonal to ∇fp(w), then applies ascent in the orthogonal component to minimize h.By construction, the orthogonal ascent step does not change fp.
  • Implementation: GSAM’s surrogate-gradient implementation adds negligible computation compared with the forward and backward passes and can be supplied to first-order optimizers.The method uses noisy gradient observations and projections to construct the update direction.
  • Gradient decomposition and ascent: The first GSAM update is identical to SAM: gradient descent on ∇fp(w) minimizes the perturbed loss.The orthogonal ascent step supplements this SAM-like descent rather than replacing it.
  • Toy example: In the toy loss surface, SGD and SAM reach sharp local minima, whereas GSAM moves toward a ridge and tends to find a flat region.The paper attributes this trajectory to the smaller surrogate gap near the smoother ridge.

5 THEORETICAL PROPERTIES OF GSAM

The paper establishes convergence for GSAM under non-convex stochastic optimization and connects surrogate-gap reduction to a tighter generalization bound than SAM.

  • Convergence: The convergence theorem assumes an L-Lipschitz-smooth non-convex function with lower bound fmin and bounded noisy gradient observations.It also specifies learning-rate and perturbation-amplitude schedules proportional to inverse square-root time.
  • Convergence: GSAM converges for non-convex stochastic optimization at rate O(log T/The stated rate matches first-order gradient optimizers such as Adam.
  • Convergence: Both the perturbed loss fp and original loss f converge in GSAM at the stated rate.
  • Generalization: Under the PAC-Bayesian analysis, minimizing the surrogate gap is expected to produce a tighter upper bound on generalization performance.The bound uses the empirical training loss and surrogate gap evaluated on the training set.
  • Generalization: GSAM provably achieves a non-trivial total decrease in surrogate gap relative to SAM, implying provably better generalization under the stated bound.

6 EXPERIMENTS

Experiments across ResNets, ViTs, and MLP-Mixers show that GSAM improves generalization over SAM and vanilla training, while surrogate-gap estimates track Hessian dominant eigenvalues. Additional studies show benefits across optimizers, augmentations, transfer tasks, and method combinations.

  • 6.1 GSAM improves test performance on various model architectures: GSAM consistently improves ImageNet generalization over SAM and vanilla training across ResNets, ViTs, and MLP-Mixers.On ViT-B/32, GSAM improves top-1 accuracy by 3.2% over SAM and 5.4% over AdamW; on Mixer-B/32, improvements are 1.2% and 11.1%, respectively.
  • 6.1 GSAM improves test performance on various model architectures: GSAM also improves out-of-distribution performance, including ImageNet-R and ImageNet-C.For Mixer-B/32, the reported improvements are 5.1% on ImageNet-R and 5.9% on ImageNet-C.
  • 6.2 GSAM finds a minimum whose Hessian has small dominant eigenvalues: Surrogate-gap estimates of the Hessian’s dominant eigenvalue match power-iteration trends, supporting the surrogate gap as a sharpness proxy.The estimated dominant eigenvalue follows decreases with α and ρ, while the surrogate gap is computed over the full training set and direct eigenvalues over a subset.
  • 6.3 Comparison with methods in the literature: Applying the GSAM ascent step improves test accuracy when combined with both SAM and ASAM, outperforming Entropy-SGD and vanilla training.The ascent step is presented as an optimizer-compatible addition, while Entropy-SGD uses an inner Langevin iteration.
  • 6.4 Additional studies: Directly minimizing fp(w) + λh(w) is empirically inferior to GSAM because the surrogate-gap gradient can conflict with perturbed-loss and loss gradients.The comparison was conducted on ViT-B/32 across searched λ and ρ values.
  • 6.4 Additional studies: GSAM maintains its advantage across data augmentations, base optimizers, and transfer-learning tasks.It outperforms SAM and vanilla training under different Mixup settings, performs best with AdamW and AdaBelief, and yields better downstream transfer performance.

7 CONCLUSION

The paper introduces the surrogate gap as an optimizable sharpness measure and GSAM as a low-overhead method that improves generalization over SAM across architectures.

  • The surrogate gap measures sharpness, is easy to compute, and is feasible to optimize directly.
  • GSAM improves generalization over SAM with negligible computation cost and is validated on diverse neural-network architectures.The paper reports theoretical convergence and better generalization, alongside empirical validation across models.

ETHICS STATEMENT

The paper develops optimization methodologies applicable to different deep neural networks and applications; ethical impact therefore depends primarily on the trained models.

  • The ethical impact is primarily determined by the specific models trained using the proposed optimization strategy.

REPRODUCIBILITY STATEMENT

The paper analyzes GSAM’s convergence and loss-gap reduction, showing how its orthogonal ascent step decreases the surrogate gap without changing the perturbed loss.

  • Lemma 3.3: At a local minimum, the surrogate gap approximates the dominant Hessian eigenvalue as σmax ≈ 2h(w)/ρ2.This connects the directly minimized loss gap to sharpness when the perturbation radius is small.
  • Loss-gap reduction: The orthogonal ascent step is guaranteed to decrease the loss gap, whereas gradient descent may increase it in practice.The analysis attributes the ascent-step decrease to its effect on the surrogate gap while preserving the perturbed-loss objective.
  • Loss-gap reduction: The loss-gap decrease depends on gradient and Hessian quantities, including a lower bound on the gradient norm and the minimum absolute Hessian eigenvalue.Under these conditions, the authors conclude that the ascent step efficiently improves generalization compared with SAM.

A.7 DISCUSSION ON COROLLARY 5.2.1

The discussion clarifies that perturbed-loss risk bounds require matching training and true perturbation amplitudes; when they differ, minimizing the surrogate gap improves the bound and generalization rationale.

  • Corollary interpretation: When ρtrain ≠ ρtrue, minimizing both h and fp is more effective for generalization than minimizing fp alone.The equality fp = C + h applies only when the training perturbation matches the ground-truth perturbation.
  • Risk bound: The risk bound is expressed as C + ρ2σ/2 + R after ignoring the higher-order O(ρ3) term under the small-radius Taylor approximation.Here C is training loss, σ is the dominant Hessian eigenvalue, and R is the regularization term.
  • Risk-bound comparison: A low perturbed training loss is insufficient for generalization when the training and true perturbation amplitudes differ; a low dominant Hessian eigenvalue is also necessary.The discussion compares solutions with the same training perturbed loss and shows why lower sharpness matters for the risk bound.
  • Perturbation mismatch: The perturbation radius increases with data variability, so broader variability between training and underlying distributions can make ρtrue exceed ρtrain.The discussion links ρ to the standard deviation of weights obtained from repeated samples of the underlying distribution.
  • Experimental settings: The experiments search optimizer, learning-rate, weight-decay, perturbation-radius, and GSAM α configurations across ResNets, ViTs, and Mixers.Downstream evaluations use ImageNet-1k pretrained models on CIFAR10, CIFAR100, Oxford-Flowers, and Oxford-IITPets.

B.3 EXPERIMENTAL SETUP WITH ABLATION STUDIES ON DATA AUGMENTATION

The experiments examine GSAM across data-augmentation strengths, neighborhood-radius settings, and ascent-step scaling, including constant and decayed radius schedules. GSAM consistently outperforms SAM and vanilla training, while the ascent step accounts for the main generalization improvement.

  • Data augmentation: The ablations vary data augmentation from Light to Medium and Strong settings, with mixup and RandAugment defining the latter two.Light uses Inception-style random flip and crop; Medium and Strong use mixup with probability 0.2 and RandAugment magnitudes 10 and 15, respectively.
  • Influence of ρ and α: Fine-tuning SAM’s neighborhood radius ρ does not achieve performance comparable to GSAM.The comparison is reported for ViT-B/32 models across different ρ values.
  • Influence of ρ and α: GSAM consistently outperforms SAM and vanilla training across tested ascent-step scaling values α.The ViT-B/32 experiments vary α and report consistent superiority over both baselines.
  • Constant ρ versus decayed ρt schedule: GSAM’s ascent step improves accuracy with both constant ρ and a decayed ρt schedule.Without the ascent step, constant and decayed schedules achieve similar performance, implicating the ascent step as the main improvement.

C.3 VISUALIZE THE TRAINING PROCESS

The training-process analysis tests a small-angle assumption and tracks the surrogate gap during optimization. Gradients remain closely aligned, while increasing α lowers the surrogate gap even as the gap rises over training steps.

  • Gradient alignment: The cosine between ∇f(wt) and ∇fp(wt) stays above 0.9 during training.This supports the analysis assumption that θt, the angle between the two gradients, is small.
  • Surrogate gap dynamics: Increasing α decreases the surrogate gap during training, indicating that GSAM’s ascent step minimizes the gap efficiently.The surrogate gap is plotted under different α values.
  • Surrogate gap dynamics: For any fixed α, the surrogate gap increases with training steps as optimization progressively approaches a local minimum while reducing training loss.The reported trend links increasing gap values with the training process’s movement toward local minima.
Loading 2203.08065v2…