PaperlayerView exact revision on arXiv

Source-linked AI summary

Sharpness-Aware Training for Free

Jiawei Du, Daquan Zhou, Jiashi Feng, Vincent Y. F. Tan, Joey Tianyi Zhou

arXiv:2205.14083v3cs.LGcs.AIcs.CV

TL;DR

Over-parameterized neural networks can have large generalization error, and SAM-like methods address sharp minima but double base-optimizer computation. SAF uses a trajectory loss based on outputs from current and past weights, while MESA offers a memory-efficient variant. The paper reports that SAF and MESA outperform SAM and its variants on ImageNet at nearly twice SAM’s training speed, with SAF adding almost no computation.

  • Problem

    Over-parameterized neural networks may have large generalization errors, while SAM-like sharpness-aware training incurs twice the base optimizer’s computational cost.

    Source: §1 · pp. 1–2

  • Method

    SAF replaces SAM’s sharpness measure with a KL-divergence trajectory loss between outputs from current and past weights, while MESA uses an exponential-moving-average network for memory efficiency.

    Source: §1 · p. 2

  • Results

    On ImageNet, SAF achieves 99.2% training speed versus SAM’s 50%, and SAF and MESA outperform SAM and its variants.

    Source: §1 · p. 2, §4.2 · p. 8

  • Takeaways & Limitations

    SAF provides sharpness-aware training with almost no additional computation, while MESA reduces memory demands for extremely large datasets at 15% additional computation.

    Source: §1 · pp. 2–3

    SAF’s memory usage grows with dataset size and may cause out-of-memory issues on extremely large datasets such as ImageNet-21K and JFT-300M.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

Modern deep neural networks (DNNs) have achieved state-of-the-art performances but are typically over-parameterized. The over-parameterization may result in undesirably large generalization error in the absence of other customized training strategies. Recently, a line of research under the name of Sharpness-Aware Minimization (SAM) has shown that minimizing a sharpness measure, which reflects the geometry of the loss landscape, can significantly reduce the generalization error. However, SAM-like methods incur a two-fold computational overhead of the given base optimizer (e.g. SGD) for approximating the sharpness measure. In this paper, we propose Sharpness-Aware Training for Free, or SAF, which mitigates the sharp landscape at almost zero additional computational cost over the base optimizer. Intuitively, SAF achieves this by avoiding sudden drops in the loss in the sharp local minima throughout the trajectory of the updates of the weights. Specifically, we suggest a novel trajectory loss, based on the KL-divergence between the outputs of DNNs with the current weights and past weights, as a replacement of the SAM's sharpness measure. This loss captures the rate of change of the training loss along the model's update trajectory. By minimizing it, SAF ensures the convergence to a flat minimum with improved generalization capabilities. Extensive empirical results show that SAF minimizes the sharpness in the same way that SAM does, yielding better results on the ImageNet dataset with essentially the same computational cost as the base optimizer.

1 Introduction

Over-parameterized neural networks can have large generalization error, while SAM seeks flatter minima at twice the base optimizer’s computational cost. SAF replaces SAM’s sharpness measure with a trajectory loss, and MESA offers a memory-efficient variant.

1 Introduction

  • Over-parameterized neural networks may have large generalization errors, motivating training strategies that address overfitting.Prior work connects flatter loss minima with better generalization.Source: §1 · p. 1
  • SAM reduces generalization error by seeking flat minima, but its sharpness approximation makes its computational cost twice that of the base optimizer.The two update steps approximate sharpness and then minimize it together with the vanilla loss.Source: §1 · pp. 1–2
  • SAF and MESA outperform SAM and its variants on ImageNet-1k while achieving almost twice SAM’s training speed; SAF adds essentially no computational overhead.The reported training speeds are 99.2% for SAF and 50% for SAM; MESA reaches 84%–85% of SGD’s speed.Source: Figure 1 · p. 2, §1 · pp. 2–3
  • SAF replaces SAM’s sharpness measure with a trajectory loss based on KL divergence between outputs from current and past weights.The method uses past network outputs as trajectory-loss targets and requires almost zero extra computation.Source: §1 · p. 2
  • MESA uses an exponential-moving-average network to provide trajectory-loss targets, reducing SAF’s memory demands for extremely large datasets at 15% additional computation.The paper identifies ImageNet-21K and JFT-300M as examples of datasets where SAF may run out of memory.Source: §1 · pp. 2–3

2 Preliminaries

The preliminaries frame training as minimizing empirical loss over a dataset and describe SAM’s alternative objective, which also penalizes sharpness in a neighborhood around the weights.

2 Preliminaries

  • Training seeks weights that minimize empirical loss over a dataset of i.i.d. samples drawn from a natural distribution.The paper uses cross-entropy loss and randomly samples fixed-size minibatches during training.Source: §2 · p. 3
  • The setup assumes the empirical loss is continuous and differentiable with bounded first-order derivative.The loss may be arbitrary, though the paper uses cross-entropy.Source: §2 · p. 3
  • SAM minimizes the vanilla loss plus a sharpness measure defined by the maximum increase in loss within a ρ-constrained neighborhood of the weights.The perturbation vector ϵ maximizes training loss within that neighborhood.Source: §2 · p. 3
  • SAM approximates the sharpness measure using the loss difference between the perturbed and original weights, with the perturbation found by a first-order Taylor approximation.The perturbation is an approximate solution to the neighborhood maximization.Source: §2 · p. 3

3 Methodology

The methodology replaces SAM’s sharpness measure with a trajectory-based proxy, then uses KL divergence to penalize changes along training. MESA adapts this approach with EMA weights to reduce memory demands and emphasize recent weights.

3 Methodology

  • SAF replaces SAM’s sharpness measure with a trajectory loss at essentially zero extra computational overhead over the base optimizer.The methodology also introduces MESA as a memory-efficient SAF variant.Source: §3 · p. 4

3.1 General Idea: Leverage the trajectory of weights to estimate the sharpness

  • Because consecutive mini-batches are usually different, direct estimation from consecutive training-loss changes is unavailable; the trajectory of standard weight updates provides a proxy for sharpness.The loss difference is described as equivalent to SAM’s sharpness measure, without additional computational overhead.Source: §3.1 · p. 5

3.2 Sharpness-Aware Training for Free (SAF)

  • SAF minimizes vanilla loss together with a KL-divergence trajectory loss, using softened outputs from ˜E epochs earlier as targets.The trajectory loss is introduced after a predefined starting epoch because early network outputs are not stable or reliable.Source: §3.2 · pp. 5–6

3.3 Memory-Efficient Sharpness-Aware Training (MESA)

  • MESA replaces stored past outputs with EMA-model outputs to construct the trajectory loss, targeting SAF’s dataset-proportional memory use and its learning-rate-linked decay in sharpness weighting.The EMA weighting places more emphasis on sharpness at current and recent weights.Source: §3.3 · p. 6

4 Experiments

Experiments compare SAF and MESA with vanilla training, SAM, and SAM variants across image-classification benchmarks. The authors report stronger results on ImageNet and CIFAR, alongside reduced sharpness and flat loss landscapes.

4.1 Experiment Setup

  • Experiments cover CIFAR-10, CIFAR-100, and ImageNet, using ResNet, Wide ResNet, PyramidNet, and Vision Transformer models against vanilla, SAM, and SAM variants.Training hyperparameters are held constant across baselines and proposed methods for fair comparison.Source: §4.1 · p. 7

4.2 Experimental Results

  • Up to 0.9%, SAF and MESA outperform SAM on ImageNet; SAF reaches 99.2% vs. SAM's 50% training speed.The table reports ImageNet classification accuracies and training speeds; the paper also says MESA runs at 84% to 85% of SGD speed.Source: Table 1 · p. 7, §4.2 · pp. 7–8

4.3 Discussion

  • SAF and MESA preserve the convergence rate, reduce SAM-measured sharpness, and yield loss landscapes reported as flat as SAM's.Sharpness decreases significantly from epoch 5, when trajectory-loss minimization begins; SAM's sharpness is lower in the second half of training.Source: Figure 4 · p. 8, §4.3 · p. 8, Figure 5 · p. 9
  • On ImageNet with ResNet-50, MESA uses 99.3% less additional memory than SAF.The authors say this reduction allows MESA to be deployed on extremely large-scale datasets.Source: Table 2 · p. 8, §4.3 · p. 8

4.2 Experimental Results

  • Across three runs, SAF and MESA outperform SAM on large-scale Wide ResNet and PyramidNet models on CIFAR-10 and CIFAR-100.Table 3 reports classification accuracies and training speeds for both CIFAR datasets.Source: §4.2 · p. 8, Table 3 · p. 9

5 Other Related Works

Prior work connects loss-landscape sharpness with generalization and develops sharpness-aware training methods. Later approaches address SAM’s computational cost or seek to improve its performance.

5 Other Related Works

  • Keskar et al. [14] linked their sharpness measure negatively with generalization ability, while Dinh et al. [4] related sharpness to Hessian spectrum and Jiang et al. [13] found a sharpness measure most correlated with generalization among 40 complexity measures.The passage situates these findings within broader studies of flat minima and generalization [4] [7] [13] [14] [17] [19] [23].Source: §5 · p. 9
  • SAM [8] and concurrent adversarial weight-perturbation methods [30] [33] advanced sharpness-aware training, but doubled the base optimizer’s computational overhead.LookSAM [20] and ESAM [6] were subsequently proposed to alleviate SAM’s computational cost.Source: §5 · p. 9
  • GSAM [36] sought to improve SAM by emphasizing gradients that minimize sharpness loss, and was reported to achieve the best generalization among post-SAM research discussed.Source: §5 · p. 9

6 Conclusion and Future Research

The paper presents trajectory loss as a low-overhead sharpness measure and proposes SAF and its memory-efficient variant MESA. It reports accuracy results on CIFAR and ImageNet, and identifies automatic training adaptation and broader task evaluation as future work.

6 Conclusion and Future Research

  • The paper introduces trajectory loss as an equivalent sharpness measure requiring almost no additional computation while preserving sharpness-aware training performance.It penalizes sharp and sudden drops in the objective function.Source: §6 · p. 10
  • SAF achieves benchmark accuracy performance on CIFAR and ImageNet while training at the speed of non-sharpness-aware methods such as SGD.Source: §6 · p. 10
  • MESA is presented as a memory-efficient SAF variant intended to make sharpness-aware training available on extremely large-scale datasets.Source: §6 · p. 10
  • Future work will adapt SAF to sense training state and alter its dynamics, evaluate more tasks such as NLP, and develop a general-purpose strategy.Source: §6 · p. 10

A.1 Derivation of Equation 8 and Equation 12

The appendix derives Equation 8 by substituting Equation 5 and considers replacing the weight trajectory with EMA weights. The resulting coefficients emphasize more recent weights.

A.1 Derivation of Equation 8 and Equation 12

  • The derivation substitutes Equation 5 into the left-hand side of Equation 8.Source: §A.1 · p. 13
  • The alternate derivation replaces trajectory Θ with WEMA, defined from successive weights and gradients, before substituting into Equation 8.Source: §A.1 · p. 13
  • The added coefficient β^(t−i) increases with i for 0 < β < 1, giving more weight to recent parameters.Source: §A.1 · p. 13

A.2 Derivation of Equation 11

The appendix derives the EMA weights at iteration t, expands them as a weighted sum, and substitutes the parameter-update expression into that sum.

A.2 Derivation of Equation 11

  • The derivation starts from the definition of the EMA weights v_t at iteration t.Source: §A.2 · p. 13
  • The EMA expansion expresses v_t using v_1 and a weighted sequence of parameters from θ_2 through θ_t.Source: §A.2 · p. 13
  • The parameter sequence is expanded through successive updates, expressing θ_t in terms of earlier parameters and gradients.Source: §A.2 · p. 13
  • The expanded parameter updates are substituted into the EMA expression to obtain the next derivation step.Source: §A.2 · p. 14

A.3 Training Details

Training hyperparameters were selected by grid search on an isolated CIFAR validation set, while the SAF schedule parameters were fixed by dataset and model family.

A.3 Training Details

  • Grid searches varied learning rate, weight decay, and batch size for SGD, SAM, SAF, and MESA, selecting settings by test accuracy.The search used an isolated validation set of CIFAR datasets; the selected settings are reported for CIFAR and ImageNet.Source: §A.3 · p. 14
  • The schedule {Estart, ˜E, τ} was set to {5, 3, 5} for CIFAR-10/100, {10, 3, 5} for ImageNet ResNets, and {150, 3, 5} for the ImageNet ViT.The settings were consistent across experiments within the CIFAR and ImageNet datasets.Source: §A.3 · p. 14
  • Tables 4 and 5 report hyperparameters for training from scratch on CIFAR-10/100 and ImageNet, respectively.The passages identify each table's scope but do not provide its cell values.Source: Table 4 · p. 15, Table 5 · p. 16

Related papers

Loading 2205.14083v3…