PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Efficient and Scalable Sharpness-Aware Minimization

Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, Yang You

arXiv:2203.02714v1cs.LGcs.AIcs.CV

TL;DR

SAM improves performance but incurs substantial overhead from two sequential gradient computations at every step. The paper proposes LookSAM, which periodically computes and reuses the inner-gradient direction, and extends it with layer-wise scaling for large-batch ViT training. LookSAM retains similar accuracy gains with first-order-like computational complexity, while Look-LayerSAM scales ViT training to a 64k batch size and reports about 8× speedup in a cited 4k-batch setting.

  • Problem

    SAM can improve accuracy but requires two sequential, non-parallelizable gradient computations per step, which doubles computational complexity compared with SGD or Adam.

  • Method

    LookSAM periodically computes SAM’s inner gradient ascent and reuses its flatness-promoting direction, while Look-LayerSAM adds layer-wise scaling for weight perturbation.

  • Results

    LookSAM achieves similar accuracy to SAM with computational complexity comparable to SGD or Adam; Look-LayerSAM scales ViT training to 64k batch size and reports about 8× speedup with a 4k batch size.

  • Takeaways & Limitations

    The proposed methods support efficient SAM-like optimization and large-batch ViT training while maintaining competitive performance.

Abstract

from arXiv · show

Recently, Sharpness-Aware Minimization (SAM), which connects the geometry of the loss landscape and generalization, has demonstrated significant performance boosts on training large-scale models such as vision transformers. However, the update rule of SAM requires two sequential (non-parallelizable) gradient computations at each step, which can double the computational overhead. In this paper, we propose a novel algorithm LookSAM - that only periodically calculates the inner gradient ascent, to significantly reduce the additional training cost of SAM. The empirical results illustrate that LookSAM achieves similar accuracy gains to SAM while being tremendously faster - it enjoys comparable computational complexity with first-order optimizers such as SGD or Adam. To further evaluate the performance and scalability of LookSAM, we incorporate a layer-wise modification and perform experiments in the large-batch training scenario, which is more prone to converge to sharp local minima. We are the first to successfully scale up the batch size when training Vision Transformers (ViTs). With a 64k batch size, we are able to train ViTs from scratch in minutes while maintaining competitive performance.

1. Introduction

SAM improves generalization by favoring flatter minima but requires two sequential gradient computations per step. LookSAM reuses a periodically computed direction, while Look-LayerSAM extends this approach to large-batch ViT training.

  • Sharp local minima are associated with degraded deep-network generalization, motivating SAM’s explicit preference for flatter regions.
  • SAM requires two sequential, non-parallelizable gradient computations at every step, increasing training overhead.
  • LookSAM periodically computes the inner gradient ascent and reuses the direction that promotes flat regions between computations.
  • Look-LayerSAM applies layer-wise weight-perturbation scaling and trains ViTs with a 64k batch size within an hour while maintaining competitive performance.
  • LookSAM achieves similar accuracy gains to SAM with computational complexity comparable to SGD or Adam.

2. Related Work

Related work links sharp minima to deep-network generalization and identifies large-batch training as both useful for distributed learning and prone to optimization and generalization challenges.

  • Sharp Local Minima: Sharp local minima can substantially influence the generalization performance of deep networks.
  • Sharp Local Minima: Prior work studies learning rate, batch size, and gradient covariance as factors influencing the minima found by SGD.
  • Large-Batch Training: Large-batch training can improve cluster utilization and accelerate training, but it introduces additional challenges.
  • Large-Batch Training: Large-batch training is prone to sharp local minima and a generalization gap because fixing epochs reduces the number of interactions.
  • Large-Batch Training: Layer-wise adaptive learning-rate methods have been proposed to scale batch sizes, including LARS for 32k-batch ResNet-50 training.

3. Method

LookSAM reduces SAM’s overhead by periodically recomputing the sharpness-aware direction and reusing its flat-region component between recomputations. Layer-wise perturbation scaling extends this approach to large-batch training, while experiments motivate the design through gradient stability and accuracy comparisons.

  • SAM overhead: SAM requires two sequential, non-parallelizable gradient computations per update, doubling computational cost relative to first-order optimizers.One gradient obtains the adversarial weight perturbation; the other computes the update at perturbed weights.
  • Naive periodic SAM: SAM-5 reduces computation but drops ViT-B-16 accuracy from 79.4% with SAM to 75.7%, compared with 74.7% for vanilla ViT.The naive strategy applies SAM only every fifth step and substantially loses SAM’s accuracy benefit.
  • Gradient decomposition: SAM’s update is decomposed into an ordinary loss-reduction component gh and a flat-region component gv, with most additional cost attributed to gv.The component gv is obtained by projecting the SAM gradient relative to the SGD direction.
  • Gradient reuse: Across five-step intervals, gv changes more slowly and more stably than gh and gs, supporting periodic recomputation of the SAM gradient.This stability motivates reusing gv while computing the ordinary SGD gradient at every step.
  • LookSAM algorithm: LookSAM recomputes the SAM gradient every k steps, projects gv, and combines the reused component with each intermediate clean-loss gradient using adaptive norm scaling.The scaling ratio matches the norms of the clean gradient and reused gv.
  • Layer-wise extension: Layer-wise SAM scales each perturbation dimension using a diagonal matrix Λ, addressing layer-dependent instability observed when scaling SAM or LookSAM to large batches.Look-LayerSAM combines this layer-wise inner maximization with LookSAM’s periodic gradient reuse.

4. Experimental Results

LookSAM matches SAM’s accuracy while reducing training cost, and its layer-wise extension supports competitive large-batch ViT training up to 64k. Results across CIFAR-100, ImageNet-1k, and large-batch settings show accuracy, efficiency, and scalability benefits.

  • CIFAR-100: LookSAM-5 matches SAM accuracy on CIFAR-100 while using much less training time across ResNet-18, ResNet-50, and WRN-28-10.Its accuracies are 80.7%, 83.3%, and 84.4% across the three models.
  • CIFAR-100: LookSAM-5 improves WRN-28-10 accuracy over SAM-5 by 0.6%, 1.0%, and 0.7% for k = 5, 10, and 20.
  • CIFAR-100: The average improvement of LookSAM-k over SAM-k increases from 0.37% to 0.53% and 0.77% as model size grows from ResNet-18 to ResNet-50 and WRN-28-10.
  • ImageNet-1k ViT: LookSAM-5 raises ViT-B-16 top-1 accuracy from 74.7% to 79.8%, whereas SAM-5 reaches 75.7%.
  • ImageNet-1k ViT: LookSAM-5 reduces ViT-B-16 training time by 2/3, from 103.1s to 68.6s, without loss in test accuracy at 79.8%.
  • Large-batch ViT: Look-LayerSAM maintains accuracy improvements of 5.6%, 5.8%, 4.4%, and 5.5% over LAMB from batch sizes 4,096 to 32,768.
  • Large-batch ViT: At 64k batch size with RandAug and Mixup, Look-LayerSAM reaches 74.9% accuracy and improves to 75.6% after Mixup.
  • Large-batch ViT: Look-LayerSAM trains ViT-B-16 in 0.7 hour with 77.1% top-1 accuracy on ImageNet-1k at 32K batch size, outperforming LAMB and SAM.

5. Conclusions

LookSAM reduces SAM’s additional computation, while Look-LayerSAM extends it to large-batch Vision Transformer training. The experiments scale batch size to 64k with accuracy above 75% and complete ViT training in 0.7 hour.

  • LookSAM reduces SAM’s additional computation and speeds up training.
  • Look-LayerSAM uses layer-wise scaling of weight perturbations for large-batch training.
  • 64k batch size achieves accuracy above 75% in Vision Transformer training.
  • 8× speedup over the 4k-batch training settings in [10] enables ViT training in 0.7 hour.

A.1. Theoretical Analysis of Projected Gradient

The appendix analyzes projected-gradient behavior through Taylor expansion, Lagrangian optimization, and the distance traveled by the gradient approximation over k steps.

  • The analysis uses Taylor expansion of the SAM loss function to derive relationships for its projected-gradient formulation.
  • The appendix analyzes the distance of gv over k steps and constrains gv to be vertical to the SGD gradient on the original weight w.
  • A Lagrangian with variables w, λ0, and λ is introduced to formulate and solve the constrained optimization problem.
  • The derivation sets partial derivatives to zero while omitting high-order terms considered trivial relative to first-order terms.

A.2. LayerSAM & LookLayerSAM

LayerSAM and Look-LayerSAM provide layer-wise gradient and perturbation updates, with Look-LayerSAM computing the SAM gradient periodically according to update frequency k.

  • LayerSAM computes a SAM gradient using layer-wise gradient norms and normalized loss gradients before updating weights.
  • Look-LayerSAM samples a minibatch, computes its gradient, and triggers the SAM-gradient computation only when t%k = 0.
  • When the periodic condition is not met, Look-LayerSAM continues with the alternative update path after the minibatch gradient computation.
  • The Look-LayerSAM update computes gv by removing the component of gs aligned with g using cos(θ).

A.3. Parameter Settings

The appendix specifies Vision Transformer architectures and training hyperparameters for vanilla and large-batch settings, including learning rate, warmup, optimizer, clipping, and epochs.

  • The Vision Transformer architectures used in the paper are listed in Table 8.
  • Vanilla ViT training parameters include learning rate, warmup, optimizer, gradient clipping, and epoch settings.
  • Large-batch ViT training uses parameter settings listed in Table 10.

A.4. Generalization bound

The section derives a high-probability generalization bound for LookSAM using PAC-Bayesian analysis and a Gaussian model of its imitated perturbations. The resulting bound relates expected test error to sharpness under perturbed weights and complexity terms involving parameter norms and logarithmic factors.

  • Theorem and proof strategy: Theorem 1 establishes a generalization bound for LookSAM with probability 1 - δ over the training set.The proof invokes a PAC-Bayesian generalization theorem for posterior and prior distributions over parameters.
  • Perturbation model: LookSAM computes the SAM gradient every k steps and reuses projected components to imitate subsequent SAM weight perturbations.The analysis represents the difference between imitated and real perturbations with ϵ0.
  • Perturbation model: The bound is formulated using the expected test error under Gaussian perturbations of the model parameters.The analysis assumes that adding Gaussian perturbations does not decrease test error and uses perturbation variance ρ′.
  • Perturbation model: When LookSAM perfectly imitates SAM by reusing the projected gradient, the perturbation variance includes the additional variance component σ0^2.The analysis defines σ′2 = σ2 + σ0^2 and notes the perfect-imitation case explicitly.
  • Final bound: The resulting bound combines sharpness, expressed through the maximum perturbed loss over ||ϵ′||p ≤ ρ′, with parameter-norm and logarithmic complexity terms.The derivation substitutes the perturbation variance into the bound after applying a concentration inequality.
Loading 2203.02714v1…