Sharpness-Aware Training for Free
Authors: Jiawei Du
Zhou Daquan
Jiashi Feng
Vincent Tan
Joey Tianyi Zhou
Venue: NeurIPS 2022 Accept (2022)
Keywords: Efficient learning
generalization
training algorithm

Abstract: 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.

Sharpness-Aware Training for Free
Jiawei Dul
Daquan Zhou², Jiashi Feng2, Vincent Y. F. Tan Joey Tianyi Zhou 1.2
Centre for Frontier Al Research (CFAR), A Singapore,
*Department of Electrical and Computer Engineering, National University of Singapore BytelDnnce
a 'Department of Mathematics, National University of Singapore
Abstract
Modern deep neural networks (DNNs) have achieved state-of ofate-of-the-ar performances but typically over- per-paarmtterized. The over paramcterizzaioonn may result in undesirably large generalization error in the absence of other customized training strategies. Recently, line of research under the name of Sharpness-Awar Mini- mization (SAM) has shown that minimizing a sharpness measure, which reflects the geometry of the loss landscape, 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 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. Specif- ically, we suggest a novel trajectory loss, based on the KL-divergence between the outputs of DNNs with the current weights and past weights, as 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 flat minimum with improved generalization capabilities. Ex- tensive empirical results show that SAF minimizes the sharpness the same way that SAM does, yielding better results on the ImageNet dataset with essentially the same computational cost as the base optimizer. Our codes are available https///gitubbooo//nnggsgDiiiiiu
significantly reduce the generalization
Introduction
Despite achieving remarkable performances in many applications, powerful neural networks [2, 5, 21, 26, 26, 27, 34, 35] are typically parameterize.. Such over-parameterizee deep neural networks require advanced training strategies to ensure that their generalization errors are appropriately small [2, 32] and the adverse effects of the overfitting are alleviated. Understanding how deep neural networks generalize is perhaps the most fundamental and yet perplexing topic in deep learning.
Numerous studies expend significant amounts of efforts to understand the generalization capabilities of deep neural networks and mitigate this problem from a variety of perspectives, such as the information perspective [18]. the model compression perspective [1, 9]. and the Bayesian perspective [22, 24], etc. The loss surface geometry perspective, in particular, has attracted lot of attention from researchers recently [4, 12-14, 28]. These studies connect the generalization gap and the sharpness of the minimum's 's loss landscape, which can be characterized by the largest eigenvalue of the Hessian matrix Vife [14] where fe represents the input-output map of the neural network. In other words, (local) minimum that is located in a flatter region tends to generalize better than one that is located in sharper one [4, 13]. The recent work [8] proposes an effective and generic training algorithm, named Sharpness-Awan Minimization (SAM), to encourage the training process converge to flat
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

minimum. SAM explicitly penalizes a sharpness measure to obtain flat minima, which has achieved state-of-the-art results in several learning tasks [2. 36].
Unfortunately, SAM's computational cost is twice that compared to the given base opti- mizer, which is typically stochastic gradient de- scent (SGD). This prohibits SAM from being de- ployed extensively in highly over-parameterized networks. Half of SAM's computational over- head is used to approximate the sharpness mea- sure in its first update step. The other half is used by SAM to minimize the sharpness mea- sure together with the vanilla loss the second update step. As shown in Figure 1, SAM sig- nificantly reduces the generalization error at the expense of double computational overhead.
Main results on ImageNet
81
RedNaa-500
.SAMM
(Ours)
R0shat-101
Reshist-152
78
78
78
77
ESAM
Vanilla (SGD)
76
400
600
800
1000
1200
1400
1600
1000
Training Speed (Images/s)
Liu al. [20] and Du et al. [6] recently ad- dressed the computation issue of SAM and proposed LookSAM [20] and Efficient SAM (ESAM) [6]. respectively. LookSAM only mini- mizes the sharpness measure once in the first of every five iterations. For the other four iterations, LookSAM reuses the gradients that minimizes the sharpness measure, which is obtained by de- composing the SAM's updated gradients in the first iteration into two orthogonal directions. As result, LookSAM saves 40% computations com- pared to SAM but unfortunately suffers from performance degradation. On the other hand, ESAM proposes two strategies to save the computational overhead in SAM's 's two updating steps. The first strategy approximates the sharpness by using fewer weights for computing, the other approximates the updated gradients by only using subset containing the instances that contribute the most to the sharpness. ESAM is reported to save up to 30% computations without degrading the performance. However, ESAM's efficiency degrades in large-scale datasets and architectures (from 30% CIFAR- 10/100 to 22% on ImageNet). LookSAM and ESAM both follow SAM's path to minimize SAM's sharpness measure, which limits the potential for further improvement of their efficiencies.
Figure 1: Accuracy training speed of SGD, SAM [8], ESAM [6], GSAM [36], and SAF. GSAM is the state-of- the-art among SAM's follow-up works. Every connected line represents a method that trains ResNet-50, ResNet- 101, and 152 models on ImageNet-Ik. SAF outperforms SAM and its variants yet requires addi- tional computational overhead.
In this paper, we aim to perform sharpness-awand training with zero additional computations and yel still retuin superior generalization performance. Specifically, we introduce novel trajectory loss to replace SAM's sharpness measure loss. This trajectory loss measures the KL-divergence between the outputs of neural networks with the current weights and those with the past weights. We propose the Sharpness-Awwe Aware training for Free (SAF) algorithm to penalize the trajectory loss for sharpness-awe awwre training. More importantly, SAF requires almost zero extra computations (SAF 0.7% v.s. SAM 100%). SAF memorizes the outputs of DNNs in the process of training as the targets of the trajectory loss. By minimizing it, SAF avoids the quick converging SAF has the potential to result in out-of-memory issue on extremely large scale datasets, such as ImageNet-21K [15] and JFT-300M [25]. We also introduce a memory-efficient variant of SAF, which is (lemory-Efficient Sharpness-Aware Training (MESA). MESA adopts a DNN whose weights are the exponential moving averages (EMA) of the trained DNN to output the targets of the trajectory loss. As a result, MESA resolves the out-of-
local sharp minimum.
issue on extremely large scale datasets, at the cost of 15% additional computations (v.s. SAM 100%). As shown in Figure 2. SAF and MESA both encourage the training to converge to flat minima similarly as SAM. Besides visualizing the loss landscape, we conduct experiments on the CIFAR-IO/I00 [16] and the ImageNet [3] datasets to verify the effectiveness of SAF. The experimental results indicate that our proposed SAF and MESA outperform SAM and its variants with almost twice the training speed: this illustrated on the ImageNet- 1k dataset in Figure 1.
In nutshell, we summarize our contributions follows.
* We propose novel trajectory loss to measure the sharpness to be used for sharpness-aware training. Requiring almost zero extra computational overhead, this trajectory loss is a better loss to quantify the sharpness compared to SAM's loss.

---

TTain nC LOCK
Trai naminn LOGS
ITam 10 LOSS
Vanilla (SGD)
SAM
SAF (Ours)
MESA (Ours)
Visualizations of loss landscapes 12 17] of the PyramidNet- I 10 model the CIFAR- 100 dataset trained with SGD, SAM, our proposed SAF, and MESA. SAF encourages the networks to converge minimum with zero additional computational overhead.
Figure 2:
We address the efficiency issue of current sharpness-awarr training, which generally incurs twice the computational overhead compared to regular training. Based proposed trajectory loss, we propose the novel SAF algorithm for improved generalization ability in this paper. SAF is demonstrated outperform SAM the ImageNet dataset, with the same computational cost as the base optimizer.
We also propose the MESA algorithm as a memory-efficien variant of SAF. MESA reduces the extra memory-usage of SAF at the cost of 15% additional computations, which allows SAF/MESA deployed efficiently (both in terms of memory and computation) on extremely large-scale datasets (e.g. ImageNet-21K [15]).
Preliminaries
Throughout the paper,
to denote neural network with weight parameters O. We given
a training dataset S that contains i.i.d. samples drawn from a natural distribution D. The training of the network is typically a non-convex optimization problem which aims to search for an optimal weight vector that satisfies
arg min Ls(Je) where Ls(so) =
Σ ((o(z.)), T
(1)
S
where l can be arbitrary loss function, and we z; to denote the pair (inputs, targets) of the i-th element the training set. In this paper, we take E to be the cross entropy loss; We use || represents lg norm; assume that Ls(fe) is continuous and differentiabl,, and its first-order derivation is bounded. In each training iteration, optimizers randomly sample a mini-batch I. C with a fixed batch size.
Sharpness-Awan Aware Minimization The conventional optimization and training focuses on minimiz- ing the empirical loss of a single weight vector 0 over the training set as stated in Equation This training paradigm is known empirical risk minimization, and tends to overfit to the training set and converges to sharp minima. Sharpness-AWe Aware Minimization (SAM) [8] aims to encourage the training to converge to flatter region in which the training losses in the neighborhood around the minimizer are lower. To achieve this, SAM proposes training scheme that solves the following min-max optimization problem:
min max Le(fotc)- cliclaSS
where p is a predefined constant that constrains the radius of the neighborhood; is the weight perturbation vector that maximizes the training loss within the p-constrained neighborhood. The objective loss function of SAM can be rewritten as the sum of the vanilla loss and the loss associated to the sharpness measure, which is the maximized change of the training loss within the p-constrained neighborhood, i.c.,
i arg min [R((f)) Ls(fo)} where Rs(fo) max La(fote) Ls(fo). cleellaSp
The sharpness measure is approximated as Rs(fo) Ls(fo+e) Ls(fo), where è is the solution to an approximated version of the maximization problem where the objective the first-order Taylor series approximation of Ls(fe+)) around 0, i.e..
VoLa(fo)
=
arg max La(fete) eileella<<p
(4)
3

---

Intuitively, SAM seeks flat minima with low variation in their training losses when the optimal weights are slightly perturbed.
Methodology
The fact that SAM's computational cost twice that compared to the base optimizer its main limitation. The additional computational overhead is used to compute the sharpness term Rs in Equation We propose a new trajectory loss as a replacement of SAM's sharpness loss Rs(fo) with essentially zero extra computational overhead over the base optimizer. Next, we introduce the Sharpness-Awan Training for Free (SAF) algorithm whose pseudocode can be found Algorithm I. We first start with recalling SAM's sharpness measure loss. Then explain the intuition for the trajectory loss as a substitute for SAM's sharpness measure loss in Section 3.1. Next, we present the complete algorithm of SAF in Section 3.2 and memory -fffcieefffiientt variant of it in Section 3.3.
Algorithm 1 Training with SAF and MESA
Input: Training set 8; A network fo with weights 0: Learning rate ry, Epochs E; Iterations T per epoch; SAF starting epoch Esaats SAF coefficients A; Temperature т; SAF hyperparameter E; EMA decay factor for MESA.
for to E' do
represents the current epoch
2:
for f. = 1 to T,, Sample a mini-batch B C S do
3:
if SAF then
4:
Record the outputs: if + fo(;), where IB
5:
Load y!
( saved in E epochs ago for each 24 E B
6:
Compute L Y(e-E))
Defined in Equation 9;
7:
else if MESA then
8:
Update EMA model weights: 24 = BU + Compute (fo, fv.)
338
9:
b Defined in Equation 13
if e > Estart then
Added the trajectory loss after Estnt epoch
11:
[ = (B) + ALE
else
13:
L - (B(f))
14:
Update the weights: / 0 Vo A flat minimum solution 8.
Output:
3.1
General Idea: Leverage the trajectory of weights to estimate the sharpness
We first rewrite the sharpness measure Rs (f of SAM based on its first-order Taylor expansion. Given vector é (Equation 4) whose norm is small, we have
Ra(fo) Is(fote) - (Jo) (a(fo) + VoTn((f)) La(fo)
(oLz(fo) VoLn(fo) (
=PIVoLL(fo)]
(5)
We remark that minimizing the sharpness loss (s( is equivalent minimizing the 4nrm the gradient L (fa), which is the same gradient used to minimize the vanilla loss Lz(fx)
The learning rate 77 in the standard training (using SGD) is typically smaller than p suggested by SAM [8]. Hence, if the mini-batch B is the same for the two consecutive iterations, the change of the training loss after the weights have been updated can be approximated as follow,
LB (ffe La La(foo (eoo(
1 = Re(fa)
(6)
0
We also remark that the change of the training loss after the weights have been updated is proportional Ru 2 Hence, minimizing the loss difference equal to minimizing the sharpness loss SAM in this case.

---

This inspires us to leverage the update of the weights in the standard training to approximate SAM's sharpness mea- sure. Regrettably though, the samples in the mini-batches By and By+1 two consecutive iterations are different with high probability, which does not allow the sharpness to be computed as in Equation 6. This is precisely the reason why SAM uses additional step compute é for approximating the sharpness. To avoid these additional computations completely, we introduce novel trajectory loss that makes use of the trajectory the weights learned in the standard training procedure to measure sharpness.
Sharp Local Minimum
Flat Local Minimum
Figure 3: For the vamilla training loss La (.fo (dashed lines), the blue arrows repre- sent the trajectory during training. Left: A sharp local minimum tends to have large trajectory loss. Right: By minimizing tra- jectory loss, SAF prevents the training from converging to sharp local minimum.
We denote the past trajectory of the weights as the itera- tions progress as θ {02, 03. ....... 02-1). Hence, Og repre- sents the weights in the 1-th iteration, which also the cur- rent model update step. We use SGD as the base optimizer to illustrate our ideas. Recall that standard SGD updates the weights as 0.+1 = = -0-0111 V LB,, (fo.). We aim to derive equivalent term of the sharpness Rž, (fa,) that can be computed without addi- tional computations. The quantity Rg, (fo.) as defined in Equation is always non-negative, and thus have
arg min (fo.) arg min RB, (fae)
arg min min Ir Ris. (fe,) Re, (fo.) ) + MA (fo,) (fa))
E ly: Rs, (fo,) Re, (fo.),
a-Uuii(())
where 0g ~ Unif(@) means that 0, is uniformly distributed in the set θ 7/ is defined as 3 cos (..). where @ is the angle between the gradients that are computed using the mini-batches B, and B, respectively. Note that for * R,,,, (fo.) Ra, (fe.) is constant with respect to the variable 0,. Therefore, when we traverse all the terms that satisfy i * 1, we obtain Equation 7. Hence, the sharpness can therefore be alternatively estimated follows:
E (r Rf. (fo:)l =0
E
[TA cos (os( LB. [ Ve, (fa:)
8 (.-Uni)(())
a,-Uni((
E 0. (2)
n: Vo, L2, V (fe.) (--Unif(e)
E
(fo.) LB, (fe
8-1 (f. LB, (fe.)]
Further details of the above derivation can be found in Appendix A.I. We remark that minimizing the loss difference L2, (fe.) LB. (ffe. is equivalent minimizing the SAM's sharpness measure Ra, (fo.). Accordingly, requiring no additional computational overhead, the training loss difference a good proxy to SAM's loss Rg, (fe,) ) to quantify and penalize the sharpness.
3.2 Sharpness-Aware Training for Free (SAF)
We claborate our proposed training algorithm SAF in this subsection. To estimate the sharpness precisely, SAF only takes the update trajectory in the past E epochs into consideration. When simultaneously minimizing the vanilla loss and the training loss difference (as in Equation 8), the second term - LB, (for) will unfortunately cancel out the vanilla loss. Therefore, we replace the cross entropy loss with the Kullback-Leible (KL) divergence loss decouple the vanilla loss. We also soften the targets of the KL divergence loss using a temperature т. Accordingly, the trajectory loss e-th epoch is defined follow
I (So. vle-E)) Bi) =
Σ
1 ((e-E) 1 KL
{ 11 : -oo(o)))
B
R.CRR ( Bi E)
where B)e-B) = = ( E B}. We remark that g
(ax):
B) the network output
of the instance æ; in E epochs ago, as illustrated in Line 4 of Algorithm 1. The outputs of each
5

---

instance 2; will be recorded, and no additional computations are required during this procedure. The trajectory loss will be deployed after a predefined epoch Sstart (Line 10 of Algorithm 1), because the outputs of the DNN are neither stable nor reliable at the first few epochs. Intuitively, the trajectory loss slows down the rate of change of the training loss to avoid convergence to sharp local minima. This is illustrated in Figure 3.
3.3 Memory-Efficient Sharpness-Aware Training (MESA)
The memory usage for recording the outputs is negligible for standard datasets such as CIFAR [16] (57 MB). However, SAF's memory usage is proportional the size of the training datasets, which may result in an out-of-memory issue on extremely large-scale datasets such as ImageNet-21k [15] and JFT-300M [25]. Another limitation of SAF that the coefficient "/ will decay at the same rate as the learning rate T7 since "/t X as shown in Equation 8. Hence. the sharpness at the current weights are quantified by smaller coefficients %t. For a SAF-like algorithm to be applicable on extremely large-scale datasets and to emphasize the sharpness of the most current or recent weights, we introduce Eefier-Efficient Sharpness-Awee moving average (EMA) weighting strategy on the weights to construct the trajectory loss. The EMA weight i the t-th iteration is updated follows
Training (MESA), which adopts an exponential
3 1 + (1 3),,
(10)
where B E (0.9, 1) is the decay coefficient of EMA. Given that 01 = O,, and Oct1 = 0: V Lz, (fo)), , the EMA weight in the 1-th iteration, can be expressed
Ut - 10- 3 Vo LB. (So.) 00 +
3 Vo, (So.).
(11)
More details of this derivation can be found in Appendix A.2. Therefore, the trajectory from 0g to it is collected in the vector WEMA = (W², Wy..... w(-1),, where w; = W;-1 - 3r "n LB, (fe.). W1 Vt, and wt 0₁. If we regard the outputs of EMA model fo, as the targets of the trajectory loss, and substitute it into Equation 8,
LB. (fo.)
Lf, (fer))
E w. WEMA)
LB, (fw.)) - Lo, (fwrti)
E
180-ty/ RB, (frr RE. (fo.)).
(12)
wy WnifUnifWWaam
More details can be found in Appendix A.I. Hence, the EMA coefficients g will place more emphasis on the sharpness of the current and most recent weights since 0 < 3 < 1. The trajectory loss of MESA is
t (So,
Σ KT. Z,EB
(
(13)
We that the difference between this and the trajectory loss L (fo, Y( B)) discussed in Sec- tion 3.2 is that the target o!! (e-E) the t-iteration has been replaced by fu, (2x) for 2; B.
4 Experiments
We verify the effectiveness of our proposed SAF and MESA algorithms in this section. We first conduct experiments to demonstrate that our proposed SAF achieves better performance compared to SAM which requires twice the training speed. SAF is shown to outperform SAM and its variants in large-scale datasets and models. The main results are summarized into Tables 1 and 3. Next, we evaluate the sharpness of SAF using the measurement proposed by SAM. We demonstrate that SAF encourages the training to converge to flat minimum. We also visualize the loss landscape of the minima converged by SGD, SAM, SAF, and MESA in Figures 2 and 5, which show that both SAF's and MESA's loss landscapes are as flat SAM.

---

4.1 Experiment Setup
Datasets We conduct experiments on the following image classification benchmark datasets: CIFAR- 10. CIFAR-100 [16], and ImageNet [3]. The 1000-class ImageNet dataset contains roughly 1.28 million training images, which is the popular benchmark for evaluating large-scale training.
Models We employ a variety of widely-used DNN architectures to evaluate the performance and training speed. We use ResNet- [11], Wide ResNet-28-10 [31], and PyramidNet-11 110 [10] for the training in CIFAR-10/100 datasets. We use ResNet [11] and Vision Transformer [5] models with
various sizes on the ImageNet dataset.
Baselines We take the vanilla (AdamW for VIT. SGD for the other models), SAM [8], ESAM [6], GSAM [36], and LookSAM [20] as the baselines. ESAM and LookSAM are SAM's follow-up works that improve efficiency. GSAM achieves the best performance among the SAM's variants. We reproduce the results of the vanilla and SAM, which matches the reported results in [2, 20, 36). And we report the cited results of baselines ESAM [6]. GSAM [36], and LookSAM [20].
Implementation details We set all the training hyperparameters to be the same for a fair comparison among the baselines and our proposed algorithms. The details of the training setting are displayed in the Appendix. We follow the settings of [2, 6, 20, 36] for the ImageNet datasets, which is different from the experimental setting of the original SAM paper [8]. The codes are implemented based on the TIMM framework [29]. The ResNets are trained with a batch size of 4096, 1.4 learning rate, 90 training epochs, and SGD optimizer (momentum=0.9) over Nvidia V- GPU cards. The ViTs trained with 300 training epochs and Adam W optimizer (81 0.9, B 0.999). We only conduct the basic data augmentation for the training on both CIFAR and ImageNet (Inception- style data augmentation). The hyperparameters of SAF and MESA are consistent among various DNNs architectures and various datasets. We set 5, E 3, Exun for all the experiments,
E {0.3, 0.8} for SAF and MESA, respectively.
Table 1: Classification accuracies and training speeds on the ImageNet dataset. The numbers in parentheses (-) indicate the ratio of the training speed w.r.t. the vanilla base optimizer's (SGD's) speed. Green indicates improvement compared to SAM, whereas red suggests degradation.
ResNet-50
ResNet-101
ImageNet
Accuracy
images/s
Accuracy
images/s
Vanilla (SGD)
76.0
1,627 (100%)
77.8
1,042 (100%)
SAM [8]
76.9
802 (49.3%)
78.6
518 (49.7%)
ESAM [6]
77.1
1,037 (63.7%)
79.1
650 (62.4%)
GSAM [36]
77.2
783 (48.1%)
78.9
503 (48.3%)
SAF (Ours)
77.8
1,612 (99.1%)
79.3
1,031 (99.0%)
MESA (Ours)
77.5
1,386 (85.2%)
79.1
888 (85.4%)
ResNet-152
VIT-S/32
ImageNet
Accuracy
images/s
Accuracy
images/s
Vanilla²
78.5
703 (100%)
68.1
5,154 (100%)
SAM [8]
79.3
351 (49.9%)
68.9
2,566 (49.8%)
LookSAM' [20]
68.8
4,273 (82.9%)
GSAM [36]
80.0
341 (48.5%)
73.8
2,469 (47.9%)
SAF (Ours)
79.9
694 (98.7%)
69.5
5,108 (99.1%)
MESA (Ours)
80.0
601 (85.5%)
69.6
4,391 (85.2%)
4.2 Experimental Results
ImageNet Our proposed SAF and MESA achieve better performance on the ImageNet dataset
compared to the other competitors. We report the best test accuracies and the average training speeds
I We report in [6]. as ESAM only release their codes in Single-GPU environment.
We report the results in [36]. but failed reproduce them using the officially released codes.
2 We use base optimizers SGD for ResNet- 152 and Adam for ViT-S/32.
'The authors of LookSAM have not released their code for either Single- or Multi- GPU environments; hence report the results in [20]. LookSAM only reports results for ViTs and there results for ResNet.
7

---

Training Loss Epochs
Sharpness vs Epochs
Vani 15G2)
SAM SAF MESA
Training loss)
Sharpmo
Epochs
Epochs
(a) Training loss vs Epochs of SAF.
(b) The SAM's sharpness measure vs epochs
Figure 4: Left: The change of the vanilla loss (exclude the trajectory loss) in cach epoch. SAF does not affect the convergence of the training. Right: The change of sharpness, which is the measurement proposed by SAM with p = 0.05. SAF and MESA decrease the sharpness measure of SAM significantly.
in Table 1. The experiment results demonstrate that SAF incurs no additional computational overhead to ensure that the training speed is the same the base Aptimizer-SSti We observe that SAF and MESA perform better in the large-scale datasets. SAF trains DNN at a douled speed than SAM (SAF 99.2% vs SAM 50%). MESA achieves training speed which is 84% to 85% that of SGD. Concerning the performance, both SAF and MESA outperform SAM up to (0.9%) on the ImageNet dataset. More importantly, SAF and MESA achieve a new state-of-the-ard results of SAM's follow-up works on the ImageNet datasets trained with ResNets.
CIFARIO and CIFAR100 We ran all the experiments three times with different random seeds for fair comparisons. We summarize the results in Table 3, in the same way as we do for the experiments on the ImageNet dataset. The training speed is consistent with the results on ImageNet. Similarly, SAF and MESA outperform SAM on large-scale models Wide Resnet and PyradmidNets.
4.3 Discussion
Memory Usage We evaluate the additional memory usage of SAF and MESA on the ImageNet dataset with the ResNet-50 model. We only present the results on the ImageNet in Table 2, because the memory usage of SAF on the CIFAR dataset is negligible (57 Mb). MESA saves 99.3% memory usage compared to SAF, which allows MESA to be deployed on extremely large-scale datasets.
Convergence Rate Intuitively, SAF minimizes the trajec- tory loss to control the rate of training loss change to obtain flat minima. A critical problem of SAF may be SAF's influ- ence on the convergence rate. We empirically show that the change of the sharpness-inducing term in SAF and MESA (compared to SAM) will not affect the convergence rate during training. Figure 4a illustrates the change of the train- ing loss in each epoch of SGD, SAM, SAF, and MESA. It shows that SAF and MESA converge at the same rate SGD and SAM.
Table 2: The additional memory used by
SAF and MESA on the ImageNet dataset.
Algorithms
Extra Memory Usage
SAF MESA
14,643 MB
98 MB
Sharpness
We empirically demonstrate that the trajectory loss can be a reliable proxy of the sharpness loss proposed by SAM. We plot the sharpness, which is the measurement of SAM in Equation in each epoch of SAM, SGD, SAF, and MESA during training. As shown in Figure 4b, SAF and MESA minimize the sharpness as SAM does throughout the entire training procedure. Both SAF's and MESA's 's sharpness decrease significantly at epoch 5, from which the trajectory loss starts to be minimized. SAM's sharpness is lower than SAF's and MESA's in the second half of the training. A plausible reason is that SAM minimizes the sharpness directly. However, SAF and MESA outperform SAM in terms of the test accuracies, as demonstrated in Table 1.
Visualization of Loss Landscapes We also demonstrate that minimizing the trajectory loss is an effective way to obtain flat minima. We visualize the loss landscape of converged minima trained by SGD, SAM. SAF, and MESA in Figures 2 and 5. We follow the method to do the plotting proposed by [17], which has also been used in [2, 6]. The x- and y-axes represent two random sampled orthogonal Gaussian perturbations. We sampled 100 x 100 points with random Gaussian perturbations for visualization. The visualized loss landscape clearly demonstrate that SAF can
converge to region as flat as SAM does.
8

---

LOS
34
Train 0g
20- 85
-25 14
00
05
Vanilla (SGD)
SAM
SAF (Ours)
MESA (Ours)
Figure 5: Loss landscapes visualization of the Wide ResNet-28- 10 model on the CIFAR- 100 dataset trained
with SGD, SAM, our proposed SAF and MESA.
Table 3:
Classification accuracies and training speed on the CIFAR-IO and CIFAR-100 datasets. The numbers
parentheses (-) indicate the ratio of the training speed w.r.t. the vanilla base optimizer's (SGD's) s) speed. Green indicates improvement compared to SAM, whereas red suggests a degradation.
CIFAR-10
CIFAR-100
ResNet-18
Accuracy
images/s
Accuracy
images/s
Vanilla (SGD)
95.61 0.02
3,289 (100%)
78.32 o.
3,314 (100%)
SAM [8]
96.50 0.88
1,657 (50.4%)
80.18: 0.8
1,690 (51.0%)
SAF (Ours)
96.37
3,213 (97.6%)
80.06 ouns
3,248 (98.0%)
MESA (Ours)
96.24 0.02
2,780 (84.5%)
79.79+ 0.09
2,793 (84.3%)
ResNet-101
Accuracy
images/s
Accuracy
images/s
Vanilla (SGD)
96.522 0.04
501 (100%)
80.68 0.16
501 (100%)
SAM [8]
97.01 0.32
246 (49.1%)
82.99 0.44
249 (49.7%)
SAF (Ours)
96.93 0.05
497 (99.2%)
82.84 0.19
497 (99.2%)
MESA (Ours)
96.90 0:23
425 (84.8%)
82.51 0.27
426 (85.0%)
Wide-28-10
Accuracy
images/s
Accuracy
images/s
Vanilla (SGD)
96.50+ 0.05
732 (100%)
31.67+a.1
739 (100%)
SAM [8]
97.07 0.11
367 (50.1%) 1%)
83.51 0.04
370 (50.0%)
SAF (Ours)
97.08
727 (99.3%)
83.81 0.0
729 (98.6%
MESA (Ours)
97.16 0.23
617 (84.3%)
83.59 0.2
625 (84.6%)
PyramidNet-111
Accuracy
images/s
Accuracy
images/s
Vanilla (SGD)
96.66
394 (100%)
81.94 0.06
401 (100%)
SAM [8]
97.25 0.15
194 (49.3%)
84.61 +: 0.06
198 (49.4%)
SAF (Ours)
97.34 0.06
391 (99.2%)
84.71 0.01
397 (99.0%)
MESA (Ours)
97.46 0.09
332 (84.3%)
84.73: 0.14
339 (84.5%)
5 Other Related Works
The first work that revealed the relation between the generalization ability and the geometry of the loss landscape (sharpness) can be traced back to [12]. Following that, many studies verified the relation between the flat minima and the generalization error [4, 7, 13, 14, 17, 19, 23]. Specifically, Keskar al. [14] proposed a sharpness measure and indicated the negative correlation between the sharpness measure and the generalization abilities. Dinh et al. [4] further argued that the sharpness measure can be related to the spectrum of the Hessian, whose eigenvalues encode the curvature information of the loss landscape. Jiang et al. [13] demonstrated that one of the sharpness-based measures is the most correlated one among 40 complexity measures by a large-scale empirical study.
SAM [8] solved the sharp minima problem by modifying training schemes to approximate and minimize certain sharpness measure. The concurrent works [30, 33] propose model adver- sarially perturb trained weights to bias the convergence. The aforementioned methods lead the way for sharpness-aware training despite the computational overhead being doubled over that of the base optimizer. Subsequently, LookSAM [20] and Efficient SAM (ESAM) [6] were proposed to alleviate the computational issue of SAM. Apart from efficiency issues, Surrogate Gap Guided Starpness-Awar Minimization (GSAM) [36] was proposed to further improve SAM's performance. GSAM places more emphasis on the gradients that minimize the sharpness loss to achieve the best
generalization performance among all the research that followed SAM.

---

Conclusion and Future Research
In this work, we introduce novel trajectory loss as an equivalent measure of sharpness, which requires almost no additional computational overhead and preserves the superior performance of sharpness-award training [6, 8, 20, 36]. In addition to deriving novel trajectory loss that penalizes sharp and sudden drops in the objective function, we propose a novel sharpness- aware algorithm SAF, which achieves impressive performances in terms of its accuracies on the benchmark CIFAR and ImageNet datasets. More importantly, SAF trains DNNs al the same speed as non-sharpness-awaae training (e.g., SGD). We also propose MESA as a memory-efficint variant of SAF to avail sharpness- aware training extremely large-scale datasets. In future research, will further enhance SAF to automatically sense the current state of training and alter the training dynamics accordingly. We will also evaluate SAF in more learning tasks (such natural language processing) and enhance it to become a general-purposs training strategy.
Acknowledgments and Disclosure of Funding
This work support by Joey Tianyi Zhou's A * SERC Central Research Fund (Use-inspired Basic Research) and the Singapore Government's 's Research, Innovation and Enterprise 2020 Plan (Advanced Manufacturing and Engineering domain) under Grant A18AIb004$
Vincent Tan acknowledges funding from a Singapore National Research Foundation (NRF) Fellowship (A-0005077-01-00) and Singapore Ministry of Education (MOE) AcRF Tier 1 Grants (A-0009042- 01-00 and A-80001 1-800018--01-00..
References
[1] Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via compression approach. In International Conference Machine Learning, pages 254-263. PMLR, 2018.
[2] Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training strong data augmentations. In International Conference on Learning Representations, 2021.
[3] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 248-255, 2009.
[4] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pages 2019-1028. PMLR, 2017.
[5] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x words: Transformers for image recognition scale. In International Conference on Learning Representations, 2020.
[6] Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent Tan. Efficient sharpness-aware minimization for improved training of neural networks. In International Conference on Learning Representations, 2021.
[7] Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proceedings of Uncertainty in Artificial Intelligence, 2017.
[8] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware min- imization for efficiently improving generalization. In International Conference on Learning Representations, 2020.
[9] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference Learning Representations, 2018.
10

---

[10] Dongyoon Han, Jiwhan Kim, and Junmo Kim. Deep pyramidal residual networks. In Pro- ceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pages 5927-5935, 2017.
[11] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEF/CVE Conference on Computer Vision and Puttern Recognition, pages 770-778, 2016.
[12] Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. In Proceedings of the 8th International Conference on Neural Information Processing Systems, pages 529-536, 1995.
[13] Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In International Conference on Learning Representations, 2019.
[14) Nitish Shirish Keskar, Jorge Nocedal, Ping Tak Peter Tang, Dheevatsa Mudigere, and Mikhail Smelyanskiy. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
[15] Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning. In European conference on computer vision, pages 491-507. Springer, 2020.
[16] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. CIFAR-10 10 and CIFAR- 100 datasets. URI: https://www. cs. toronto. edu/kriz/iffar html, 6(1):1, 2009.
[17] Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Proceedings of the 32nd International Conference on Neural Information Processing Systems, 31, 2018.
[18] Tengyuan Liang, Tomaso Poggio, Alexander Rakhlin, and James Stokes. Fisher- Rao metric. geometry, and complexity of neural networks. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 888-896. PMLR, 2019.
[19] Chen Liu, Mathieu Salzmann, Tao Lin, Ryota Tomioka, and Sabine Süsstrunk.. On the loss land- scape of adversarial training: Identifying challenges and how to overcome them. Proceedings of the 34th International Conference Neural Information Processing Systems, 33:21476-21487, 2020.
[20] Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Patlern Recognition, pages 12360-12370, 2022.
[21] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 10012-10022, 2021.
[22] David A McAllester. PAC-Bayesian model averaging. In Proceedings of the Twelfth Annual Conference on Computational Learning Theory, pages 164-170, 1999.
[23] Seyed-Mohsen Moosavi-Dezfooli Alhussein Fawzi, Jonathan Uesato, and Pascal Frossard. Robustness via curvature regularization, and vice versa. Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pages 9078-9086, 2019.
[24] Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A pac-bayesian approach to spertrally-normalizeed margin bounds for neural networks. In International Conference on Learning Representations, 2018.
[25] Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable effectiveness of data in deep learning era. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 843-852, 2017.
11

---

[26] Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner. Jessica Yung. Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mlp-mixer: An all-mip architecture for vision. Proceedings of the 35th International Conference on Neural Information Processing Systems, 34, 2021.
[27] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, pages 10347-10357. PMLR, 2021.
[28] Yusuke Tsuzuku, Issei Sato, and Masashi Sugiyama. Normalized flat minima: Exploring scale invariant definition of flat minima for neural networks using pac-bayesian analysis. In International Conference on Machine Learning, pages 9636-9647. PMLR, 2020.
[29] Ross Wightman. Pytorch image models.
https://github. cow/ com/rwightman
pytorch- mmm.-mmmmess.ss 2019.
[30] Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization. Proceedings of the 34th International Conference Neural Information Processing Systems, 33:2958-2969, 2020.
[31] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. British Machine Vision Conference 2016. British Machine Vision Association, 2016.
[32] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107-
115, 2021.
[33] Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. In Proceedings of the IEEE/CVF Conference Computer Vision and Puttern Recognition, pages 8156-8165, 2021.
[34] Daquan Zhou, Bingyi Kang, Xiaojie Jin, Linjie Yang. Xiaochen Lian, Zihang Jiang, Qibin Hou, and Jiashi Feng. Vie: Towards deeper vision transformer. CoRR, abs/2103.1 1186, 2021.
[35] Daquan Zhou, Zhiding Yu, Enze Xie, Chaowei Xiao, Animashree Anandkumar, Jiashi Feng, and Jose M Alvarez. Understanding the robustness in vision transformers. In International Conference on Machine Learning, pages 27378-27394, PMLR, 2022.
[36] Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha C Dvornek. James Duncan, Ting Liu, et al. Surrogate gap minimization improves sharpness-aware training. In International Conference on Learning Representations, 2021.
12

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? Yes
(b) Did you describe the limitations of your work? Yes, in Section 3.3.
(c) Did you discuss any potential negative societal impacts of your work? N.A. Our research work does not have any potential negative societal impacts.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? Yes
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? Yes
(b) Did you include complete proofs of all theoretical results? Yes
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? Yes, in Appendix.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? Yes, in Section and Appendix.
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? N.A.
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? Yes, in Section 4.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? Yes, in Section
(b) Did you mention the license of the assets? N.A.
(c) Did you include any new assets either in the supplemental material URL? N.A.
(d) Did you discuss whether and how consent was obtained from people whose data you't 're using/curating? N.A.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? N.A.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? N.A.
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? N.A.
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? N.A.
13


Reviews:
[{"review_id": "tXCVWY5tWaH", "reviewer": "Reviewer_3mAW", "summary": "This paper proposes two methods for sharpness aware minimization (SAM) without the cost of SAM, as SAM doubles training cost. The methods are 1) SAF which adds a distillation loss from the predictions of prior epochs and 2) MESA which adds a distillation loss from the exponential moving average of the weights. For resnets on imagenet, SAF and MESA match or outperform other methods including SAM, ESAM, and Vanilla SGD.", "questions": "I understand that there is not space to answer all of these questions.\n- Why is the ResNet baseline for SAM much lower than that reported in the SAM paper? \n- Why is MESA only 15% more expensive than standard training\n- The best ResNet50 in the paper is < 78%. However, modern day ResNets (e.g., from https://arxiv.org/pdf/2110.00476.pdf in the popular timm library, which is used in this paper) can get ~80%. Are the methods in this paper additive with more modern approaches.\n- Why is GSAM much better than all other methods for the ViT model (minor).", "limitations": "The paper claims to address this in Section 3.3 but a revision could benefit from an explicit limitations section (minor).", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper tackles and interesting and practical problem, which is how to make training better without the added cost of methods like SAM.\n- For ResNets on ImageNet, the empirical performance when compared to baselines is good. Moreover, this is without the cost.\n\nWeaknesses:\n- I do not understand how MESA is only 15% more expensive than standard training. If I am understanding correctly, MESA requires a forward pass through both the EMA weights and the standard weights. Doesn't this require two forward passes where there was previously only one?\n- It seems like free should not be in the name of SAF because there is an additional memory overhead (minor)\n- I did not understand very well how SAF and MESA relate to SAM. In particular, Equation 8 is not well motivated and I did not see where the KL came from.\n- I have a few concerns with the ResNet baselines: 1) The ResNet baseline for SAM is below that reported in the SAM paper -- could this perhaps be because perturbations are synced? 2) Modern ResNets, e.g. with techniques from the timm library which is the library used in this paper can get to 80% accuracy (https://arxiv.org/pdf/2110.00476.pdf). Are the methods presented in this paper additive with these modern approaches?\n- ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ahZOHFzJn98", "reviewer": "Reviewer_RMT8", "summary": "This paper proposes efficient implementations of sharpness-aware training. Specifically, this paper proposed Sharpness-Aware Training for Free (SAF) which requires cached prediction for previous checkpoints, and Memory-Efficient Sharpness-Aware (MESA) training which requires two forward propagation and one back propagation for each optimization step to achieve sharpness-aware training. While these approaches require relatively light computation costs compared to the original SAM, they achieved competitive accuracy in various tasks including CIFAR-10/100 and ImageNet-1k.", "questions": "Q1. In Fig. 2 and Fig. 5, the authors provide visualization of the loss landscape for various sharpness-aware training. Although the results showed that SAF/MESA found flatter local minima than SAM, one cannot find the clear intuition behind these empirical results. Also, while these visualizations tell us that MESA found flatter local minima than SAF for WRN28-10/CIFAR-100, the test accuracy of SAF is better than MESA. Therefore, it can be interpreted as flatter local minima considered in this paper would not result in better test accuracy in practice. Authors should clarify this paradox to remove the confusion of readers.\n\nQ2. As mentioned in the original SAM paper, the SAM optimizer can find local minima with better test accuracy with longer training (See Table 2 of [1]). Can SAF and MESA also utilize the longer training without overfitting? I think this property would be important to the scalability of SAF & MESA.", "limitations": "While this paper provides a novel interpretation of sharpness by leveraging the trajectory. several points are needed to improve the clarity and significance of the paper.\n\n1. It would be recommended to provide the detailed derivations of L163-L167. While the intuition behind these lines is clear, further clarity can be attained by providing the entire derivation of equations to readers.\n\n2. It should be clarified why SAM accuracy in Table 1 is severely lower than the original SAM paper [1]. Since they did not report any confidence interval in Table 1, this accuracy gap can change the interpretation of experimental results in Section 4.1. \n\n3. It would be recommended to add the results of SAF/MESA for longer training and stronger augmentation. Since many SOTA models in computer vision [3,4] exploit longer training and stronger augmentation, the scalability of SAF/MESA should be validated with those settings.\n\n[3] https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives\n\n[4] Liu, Zhuang, et al. \"A convnet for the 2020s.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.\n\n** After discussion phase ** \n\nThanks to the clarification of the authors, I resolved my misunderstanding of this paper. However, I think these discussion points should be reflected in the main paper: Their key trick using $\\arg \\min$ is not specified in their derivation of the main paper. Also, further clarification on differences between main references (e.g., label smoothing, number of epochs) should be clarified in their main paper since the gap between their performances is sufficiently large to change the interpretation of experimental results.\n", "rating": 3, "confidence": 4, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "* Originality: This paper provides a novel idea to estimate the sharpness by leveraging the trajectory of weights (Eq. 7). Based on this idea, they proposed time-efficient (SAF) and memory-efficient (MESA) implementations of sharpness-aware training [+].\n\n* Clarity: While Fig. 3 is intuitive to understand the key idea of this paper [+], it is still unclear to me that the equation in L164 and the second approximation in eq 7. Also, it is unclear that \\Phi_t is defined on what parameter of the optimization step [-].\n\n* Significance: I found that there exists a gap between reported SAM accuracies in previous studies and this paper. In Table 1 of this paper, the authors reported SAM accuracy as 76.9/78.6/79.3 for ResNet-50/101/152. However, the original SAM paper reported their accuracy as 77.5/79.8/80.8 with 100 epoch training (See Table 2 of [1]). I think this gap is critical since the accuracy of SAF 77.8/79.3/80.0, the minor improvements only exist for ResNet-50, and MESA 77.5/79.1/80.0, there are no improvements for SAM. Also, the ViT-S/32 results in Table 1 also make some confusion: In Table 2 of [2], they reported the accuracy of ViT-S/32-SAM as 70.5, which is severely better performing than the 68.9 in this paper. The authors should clarify this for a fair comparison.\n\n[1] Foret, Pierre, et al. \"Sharpness-aware minimization for efficiently improving generalization.\" arXiv preprint arXiv:2010.01412 (2020).\n\n[2] Chen, Xiangning, Cho-Jui Hsieh, and Boqing Gong. \"When vision transformers outperform ResNets without pre-training or strong data augmentations.\" arXiv preprint arXiv:2106.01548 (2021).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "EBEODdjDYMU", "reviewer": "Reviewer_nBtT", "summary": "This paper focuses on the computation cost perspective of Sharpness-Aware Optimization (SAM), which has 2x cost compared with SGD. \nThe first proposed algorithm, SAM for Free (SAF), derives an alternative/surrogate optimization objective which can uses the past optimization trajectory to estimate the sharpness loss. The SAF algorithm has almost the same computation cost as SGD, while inducing some extra memory usage. To mitigate the memory cost, the authors propose the second algorithm, Memory-Efficient Sharpness-Aware Training (MESA), which reduces the memory usage significantly at the tradeoff of 15% more computation cost. The empirical results of SAF and MESA are pretty good compared with SAM and SGD on popular benchmarks (CIFAR and ImageNet), indicating that the the proposed algorithms indeed reduce the computation costs of SAM while maintaining the performance.", "questions": "1. Is the reduction of computation cost from 2x to 1x SGD really meaningful? \nTheoretically speaking, 1x and 2x costs are about the same. Empirically, I know there are many factors that affecting the computation cost of training, and I'm not sure if 50% cost reduction is very meaningful practically. \n\n2.  See Weakness-1 above. That's one of my questions.", "limitations": "See the Questions section above.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The SAF algorithm is proposed in a principled way: The authors studied the SAM loss carefully, and obtained some interesting observations (Sec. 3.1), naturally leading to an elegant algorithm\n\n2. The authors indeed make the algorithm practical: Even though the general idea (Sec. 3.1) is elegant, it may still have some limitations when being used as a practical optimizer. The authors tackled some numerical challenge in Sec. 3.2, and moreover, they mitigated the memory cost issues by using the trick of exponential moving average (EMA), leading to the memory-efficient version (MESA). I really appreciate the efforts the authors spent on making the algorithm practically useful.\n\n3. The empirical results are good: The proposed algorithms, SAF and MESA, have slightly better performance than SAM with ~1/2 computation cost. Also, the extra memory cost of MESA is indeed negligible compared with SAF.\n\nWeakness: \n\n1. Is the replacement of cross-entropy with KL (in Sec. 3.2) really necessary? I think cross-entropy is still the best choice for classification, so I guess the replacement may have some negative impacts. I think it's better to have an ablation study on this replacement.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "3qx2-s-ERgq", "reviewer": "Reviewer_Byfm", "summary": "This paper proposes a sharpness-aware minimization method that adds almost no computational overhead. The authors propose to use the training loss differences (more precisely, the KL divergence of output predictions) to proxy the sharpness term. The advantage of using this new metric would be that it only requires additional memory but no additional computations. They also propose an adjusted method that addresses the memory problem of the first method. The effectiveness of the proposed methods is then shown experimentally for Imagenet and CIFAR datasets.", "questions": "1- It is mentioned that before epoch E_start the outputs are not stable. What would happen if from the first epoch this method is used? Does this instability affect the performance?\n\n2- How should the hyper-parameters (tau, E, and lambda) be selected for your method? Have you studied various hyper-parameters? It is mentioned that a fixed value is used for these hyper-parameters.\n\n3- In Figure 4b, why does MESA have higher sharpness at the beginning of training but lower sharpness at the end of training compared to SAF? Also, from the visualization, it appears that MESA has the flattest landscape.\n", "limitations": "The limitation of the method with regards to memory is mentioned and properly addressed.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1- The paper is very well-written. \n\n2- The limitation of the initial method in terms of memory has been properly addressed in the follow-up method.\n\n3- The method is simple and efficient to use.\n\n4- The experiments use state-of-the-art settings.\n\nWeaknesses:\n\n1- The method doesn't always outperform SAM in terms of accuracy. This is not so much a weakness because the method is mainly intended to improve computational overhead.\n\n2- The experiments are only for image-classification tasks.\n\n3- The method introduces a few hyper-parameters.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept