Sharpness-Aware Minimization for Efficiently Improving Generalization — part 3/4

Cleaned-text code points [31309, 47257). Bibliography omitted here; source.md includes it.

We thank our colleagues at Google --- Atish Agarwala, Xavier Garcia, Dustin Tran, Yiding Jiang, Basil Mustafa, Samy Bengio --- for their feedback and insightful discussions. We also thank the JAX and FLAX teams for going above and beyond to support our implementation. We are grateful to Sven Gowal for his help in replicating EfficientNet using JAX, and Justin Gilmer for his implementation of the Lanczos algorithm 8 8 8 https://github.com/google/spectral-density used to generate the Hessian spectra. We thank Niru Maheswaranathan for his matplotlib mastery. We also thank David Samuel for providing a PyTorch implementation of SAM 9 9 9 https://github.com/davda54/sam .

## Appendix A Appendix

### A.1 PAC Bayesian Generalization Bound

Below, we state a generalization bound based on sharpness.

###### Theorem 2 .

For any ρ > 0 \rho>0 and any distribution 𝒟 \mathscr{D} , with probability 1 − δ 1-\delta over the choice of the training set 𝒮 ∼ 𝒟 \mathcal{S}\sim\mathscr{D} , L 𝒟 ​ ( 𝒘 ) ≤ max ‖ ϵ ‖ 2 ≤ ρ ⁡ L 𝒮 ​ ( 𝒘 + ϵ ) + k ​ log ⁡ ( 1 + ‖ 𝒘 ‖ 2 2 ρ 2 ​ ( 1 + log ⁡ ( n ) k ) 2 ) + 4 ​ log ⁡ n δ + O ~ ​ ( 1 ) n − 1 L_{\mathscr{D}}(\bm{w})\leq\max_{\|\bm{\epsilon}\|_{2}\leq\rho}L_{\mathcal{S}}(\bm{w}+\bm{\epsilon})+\sqrt{\frac{k\log\left(1+\frac{\|\bm{w}\|_{2}^{2}}{\rho^{2}}\left(1+\sqrt{\frac{\log(n)}{k}}\right)^{2}\right)+4\log\frac{n}{\delta}+\tilde{O}(1)}{n-1}} (4) where n = | 𝒮 | n=|\mathcal{S}| , k k is the number of parameters and we assumed L 𝒟 ​ ( 𝐰 ) ≤ 𝔼 ϵ i ∼ 𝒩 ⁡ ( 0 , ρ ) ​ [ L 𝒟 ​ ( 𝐰 + ϵ ) ] L_{\mathscr{D}}(\bm{w})\leq\mathbb{E}_{\epsilon_{i}\sim{\mathcal{N}}(0,\rho)}[L_{\mathscr{D}}(\bm{w}+\bm{\epsilon})] .

The condition L 𝒟 ​ ( 𝒘 ) ≤ 𝔼 ϵ i ∼ 𝒩 ⁡ ( 0 , ρ ) ​ [ L 𝒟 ​ ( 𝒘 + ϵ ) ] L_{\mathscr{D}}(\bm{w})\leq\mathbb{E}_{\epsilon_{i}\sim{\mathcal{N}}(0,\rho)}[L_{\mathscr{D}}(\bm{w}+\bm{\epsilon})] means that adding Gaussian perturbation should not decrease the test error. This is expected to hold in practice for the final solution but does not necessarily hold for any 𝒘 \bm{w} .

###### Proof.

First, note that the right hand side of the bound in the theorem statement is lower bounded by k ​ log ⁡ ( 1 + ‖ 𝒘 ‖ 2 2 / ρ 2 ) / ( 4 ​ n ) \sqrt{k\log(1+\|\bm{w}\|_{2}^{2}/\rho^{2})/(4n)} which is greater than 1 when ‖ 𝒘 ‖ 2 2 > ρ 2 ​ ( exp ⁡ ( 4 ​ n / k ) − 1 ) \|\bm{w}\|_{2}^{2}>\rho^{2}(\exp(4n/k)-1) . In that case, the right hand side becomes greater than 1 in which case the inequality holds trivially. Therefore, in the rest of the proof, we only consider the case when ‖ 𝒘 ‖ 2 2 ≤ ρ 2 ​ ( exp ⁡ ( 4 ​ n / k ) − 1 ) \|\bm{w}\|_{2}^{2}\leq\rho^{2}(\exp(4n/k)-1) .

The proof technique we use here is inspired from Chatterji et al. (2020) . Using PAC-Bayesian generalization bound McAllester (1999) and following Dziugaite & Roy (2017) , the following generalization bound holds for any prior 𝒫 \mathscr{P} over parameters with probability 1 − δ 1-\delta over the choice of the training set 𝒮 \mathcal{S} , for any posterior 𝒬 \mathscr{Q} over parameters: 𝔼 𝒘 ∼ 𝒬 ​ [ L 𝒟 ​ ( 𝒘 ) ] ≤ 𝔼 𝒘 ∼ 𝒬 ​ [ L 𝒮 ​ ( 𝒘 ) ] + K L ( 𝒬 | | 𝒫 ) + log n δ 2 ​ ( n − 1 ) \mathbb{E}_{\bm{w}\sim\mathscr{Q}}[L_{\mathscr{D}}(\bm{w})]\leq\mathbb{E}_{\bm{w}\sim\mathscr{Q}}[L_{\mathcal{S}}(\bm{w})]+\sqrt{\frac{KL(\mathscr{Q}||\mathscr{P})+\log\frac{n}{\delta}}{2(n-1)}} (5) Moreover, if 𝒫 = 𝒩 ⁡ ( 𝝁 P , σ P 2 ​ 𝑰 ) \mathscr{P}={\mathcal{N}}(\bm{\mu}_{P},\sigma_{P}^{2}\bm{I}) and 𝒬 = 𝒩 ⁡ ( 𝝁 Q , σ Q 2 ​ 𝑰 ) \mathscr{Q}={\mathcal{N}}(\bm{\mu}_{Q},\sigma_{Q}^{2}\bm{I}) , then the KL divergence can be written as follows: K L ( 𝒫 | | 𝒬 ) = 1 2 [ k ​ σ Q 2 + ‖ 𝝁 P − 𝝁 Q ‖ 2 2 σ P 2 − k + k log ( σ P 2 σ Q 2 ) ] KL(\mathscr{P}||\mathscr{Q})=\frac{1}{2}\bigg[\frac{k\sigma_{Q}^{2}+\|\bm{\mu}_{P}-\bm{\mu}_{Q}\|_{2}^{2}}{\sigma_{P}^{2}}-k+k\log\left(\frac{\sigma_{P}^{2}}{\sigma_{Q}^{2}}\right)\bigg] (6) Given a posterior standard deviation σ Q \sigma_{Q} , one could choose a prior standard deviation σ P \sigma_{P} to minimize the above KL divergence and hence the generalization bound by taking the derivative 10 10 10 Despite the nonconvexity of the function here in σ P 2 \sigma_{P}^{2} , it has a unique stationary point which happens to be its minimizer. of the above KL with respect to σ P \sigma_{P} and setting it to zero. We would then have σ P ∗ 2 = σ Q 2 + ‖ 𝝁 P − 𝝁 Q ‖ 2 2 / k {\sigma^{*}_{P}}^{2}=\sigma_{Q}^{2}+\|\bm{\mu}_{P}-\bm{\mu}_{Q}\|_{2}^{2}/k . However, since σ P \sigma_{P} should be chosen before observing the training data 𝒮 \mathcal{S} and 𝝁 Q \bm{\mu}_{Q} , σ Q \sigma_{Q} could depend on 𝒮 \mathcal{S} , we are not allowed to optimize σ P \sigma_{P} in this way. Instead, one can have a set of predefined values for σ P \sigma_{P} and pick the best one in that set. See Langford & Caruana (2002) for the discussion around this technique. Given fixed a , b > 0 a,b>0 , let T = { c ​ exp ⁡ ( ( 1 − j ) / k ) | j ∈ ℕ } T=\{c\exp((1-j)/k)|j\in\mathbb{N}\} be that predefined set of values for σ P 2 \sigma_{P}^{2} . If for any j ∈ ℕ j\in\mathbb{N} , the above PAC-Bayesian bound holds for σ P 2 = c ​ exp ⁡ ( ( 1 − j ) / k ) \sigma_{P}^{2}=c\exp((1-j)/k) with probability 1 − δ j 1-\delta_{j} with δ j = 6 ​ δ π 2 ​ j 2 \delta_{j}=\frac{6\delta}{\pi^{2}j^{2}} , then by the union bound, all above bounds hold simultaneously with probability at least 1 − ∑ j = 1 ∞ 6 ​ δ π 2 ​ j 2 = 1 − δ 1-\sum_{j=1}^{\infty}\frac{6\delta}{\pi^{2}j^{2}}=1-\delta .

Let σ Q = ρ \sigma_{Q}=\rho , 𝝁 Q = 𝒘 \bm{\mu}_{Q}=\bm{w} and 𝝁 P = 𝟎 \bm{\mu}_{P}=\bm{0} . Therefore, we have: σ Q 2 + ‖ 𝝁 P − 𝝁 Q ‖ 2 2 / k ≤ ρ 2 + ‖ 𝒘 ‖ 2 2 / k ≤ ρ 2 ​ ( 1 + exp ⁡ ( 4 ​ n / k ) ) \sigma_{Q}^{2}+\|\bm{\mu}_{P}-\bm{\mu}_{Q}\|_{2}^{2}/k\leq\rho^{2}+\|\bm{w}\|_{2}^{2}/k\leq\rho^{2}(1+\exp(4n/k)) (7) We now consider the bound that corresponds to j = ⌊ 1 − k ​ log ⁡ ( ( ρ 2 + ‖ 𝒘 ‖ 2 2 / k ) / c ) ⌋ j=\lfloor 1-k\log((\rho^{2}+\|\bm{w}\|_{2}^{2}/k)/c)\rfloor . We can ensure that j ∈ ℕ j\in\mathbb{N} using inequality equation 7 and by setting c = ρ 2 ​ ( 1 + exp ⁡ ( 4 ​ n / k ) ) c=\rho^{2}(1+\exp(4n/k)) . Furthermore, for σ P 2 = c ​ exp ⁡ ( ( 1 − j ) / k ) \sigma_{P}^{2}=c\exp((1-j)/k) , we have: ρ 2 + ‖ 𝒘 ‖ 2 2 / k ≤ σ P 2 ≤ exp ⁡ ( 1 / k ) ​ ( ρ 2 + ‖ 𝒘 ‖ 2 2 / k ) \rho^{2}+\|\bm{w}\|_{2}^{2}/k\leq\sigma_{P}^{2}\leq\exp(1/k)\left(\rho^{2}+\|\bm{w}\|_{2}^{2}/k\right) (8)

Therefore, using the above value for σ P \sigma_{P} , KL divergence can be bounded as follows:

K L ( 𝒫 | | 𝒬 ) \displaystyle KL(\mathscr{P}||\mathscr{Q}) = 1 2 ​ [ k ​ σ Q 2 + ‖ 𝝁 P − 𝝁 Q ‖ 2 2 σ P 2 − k + k ​ log ⁡ ( σ P 2 σ Q 2 ) ] \displaystyle=\frac{1}{2}\bigg[\frac{k\sigma_{Q}^{2}+\|\bm{\mu}_{P}-\bm{\mu}_{Q}\|_{2}^{2}}{\sigma_{P}^{2}}-k+k\log\left(\frac{\sigma_{P}^{2}}{\sigma_{Q}^{2}}\right)\bigg] (9) ≤ 1 2 ​ [ k ⁡ ( ρ 2 + ‖ 𝒘 ‖ 2 2 / k ) ρ 2 + ‖ 𝒘 ‖ 2 2 / k − k + k ​ log ⁡ ( exp ⁡ ( 1 / k ) ​ ( ρ 2 + ‖ 𝒘 ‖ 2 2 / k ) ρ 2 ) ] \displaystyle\leq\frac{1}{2}\bigg[\frac{k(\rho^{2}+\|\bm{w}\|_{2}^{2}/k)}{\rho^{2}+\|\bm{w}\|_{2}^{2}/k}-k+k\log\left(\frac{\exp(1/k)\left(\rho^{2}+\|\bm{w}\|_{2}^{2}/k\right)}{\rho^{2}}\right)\bigg] (10) = 1 2 ​ [ k ​ log ⁡ ( exp ⁡ ( 1 / k ) ​ ( ρ 2 + ‖ 𝒘 ‖ 2 2 / k ) ρ 2 ) ] \displaystyle=\frac{1}{2}\bigg[k\log\left(\frac{\exp(1/k)\left(\rho^{2}+\|\bm{w}\|_{2}^{2}/k\right)}{\rho^{2}}\right)\bigg] (11) = 1 2 ​ [ 1 + k ​ log ⁡ ( 1 + ‖ 𝒘 ‖ 2 2 k ​ σ Q 2 ) ] \displaystyle=\frac{1}{2}\bigg[1+k\log\left(1+\frac{\|\bm{w}\|_{2}^{2}}{k\sigma_{Q}^{2}}\right)\bigg] (12) Given the bound that corresponds to j j holds with probability 1 − δ j 1-\delta_{j} for δ j = 6 ​ δ π 2 ​ j 2 \delta_{j}=\frac{6\delta}{\pi^{2}j^{2}} , the log term in the bound can be written as: log ⁡ n δ j \displaystyle\log\frac{n}{\delta_{j}} = log ⁡ n δ + log ⁡ π 2 ​ j 2 6 \displaystyle=\log\frac{n}{\delta}+\log\frac{\pi^{2}j^{2}}{6} ≤ log ⁡ n δ + log ⁡ π 2 ​ k 2 ​ log 2 ⁡ ( c / ( ρ 2 + ‖ 𝒘 ‖ 2 2 / k ) ) 6 \displaystyle\leq\log\frac{n}{\delta}+\log\frac{\pi^{2}k^{2}\log^{2}(c/(\rho^{2}+\|\bm{w}\|_{2}^{2}/k))}{6} ≤ log ⁡ n δ + log ⁡ π 2 ​ k 2 ​ log 2 ⁡ ( c / ρ 2 ) 6 \displaystyle\leq\log\frac{n}{\delta}+\log\frac{\pi^{2}k^{2}\log^{2}(c/\rho^{2})}{6} ≤ log ⁡ n δ + log ⁡ π 2 ​ k 2 ​ log 2 ⁡ ( 1 + exp ⁡ ( 4 ​ n / k ) ) 6 \displaystyle\leq\log\frac{n}{\delta}+\log\frac{\pi^{2}k^{2}\log^{2}(1+\exp(4n/k))}{6} ≤ log ⁡ n δ + log ⁡ π 2 ​ k 2 ​ ( 2 + 4 ​ n / k ) 2 6 \displaystyle\leq\log\frac{n}{\delta}+\log\frac{\pi^{2}k^{2}(2+4n/k)^{2}}{6} ≤ log ⁡ n δ + 2 ​ log ⁡ ( 6 ​ n + 3 ​ k ) \displaystyle\leq\log\frac{n}{\delta}+2\log\left(6n+3k\right) Therefore, the generalization bound can be written as follows: 𝔼 ϵ i ∼ 𝒩 ⁡ ( 0 , σ ) ​ [ L 𝒟 ​ ( 𝒘 + ϵ ) ] ≤ 𝔼 ϵ i ∼ 𝒩 ⁡ ( 0 , σ ) ​ [ L 𝒮 ​ ( 𝒘 + ϵ ) ] + 1 4 ​ k ​ log ⁡ ( 1 + ‖ 𝒘 ‖ 2 2 k ​ σ 2 ) + 1 4 + log ⁡ n δ + 2 ​ log ⁡ ( 6 ​ n + 3 ​ k ) n − 1 \mathbb{E}_{\epsilon_{i}\sim{\mathcal{N}}(0,\sigma)}[L_{\mathscr{D}}(\bm{w}+\bm{\epsilon})]\leq\mathbb{E}_{\epsilon_{i}\sim{\mathcal{N}}(0,\sigma)}[L_{\mathcal{S}}(\bm{w}+\bm{\epsilon})]+\sqrt{\frac{\frac{1}{4}k\log\left(1+\frac{\|\bm{w}\|_{2}^{2}}{k\sigma^{2}}\right)+\frac{1}{4}+\log\frac{n}{\delta}+2\log\left(6n+3k\right)}{n-1}} (13) In the above bound, we have ϵ i ∼ 𝒩 ⁡ ( 0 , σ ) \epsilon_{i}\sim{\mathcal{N}}(0,\sigma) . Therefore, ‖ ϵ ‖ 2 2 \|\bm{\epsilon}\|_{2}^{2} has chi-square distribution and by Lemma 1 in Laurent & Massart (2000) , we have that for any positive t t : P ⁡ ( ‖ ϵ ‖ 2 2 − k ​ σ 2 ≥ 2 ​ σ 2 ​ k ​ t + 2 ​ t ​ σ 2 ) ≤ exp ⁡ ( − t ) P(\|\bm{\epsilon}\|_{2}^{2}-k\sigma^{2}\geq 2\sigma^{2}\sqrt{kt}+2t\sigma^{2})\leq\exp(-t) (14) Therefore, with probability 1 − 1 / n 1-1/\sqrt{n} we have that: ‖ ϵ ‖ 2 2 ≤ σ 2 ​ ( 2 ​ ln ⁡ ( n ) + k + 2 ​ k ​ ln ⁡ ( n ) ) ≤ σ 2 ​ k ​ ( 1 + ln ⁡ ( n ) k ) 2 ≤ ρ 2 \|\bm{\epsilon}\|_{2}^{2}\leq\sigma^{2}(2\ln(\sqrt{n})+k+2\sqrt{k\ln(\sqrt{n})})\leq\sigma^{2}k\left(1+\sqrt{\frac{\ln(n)}{k}}\right)^{2}\leq\rho^{2} Substituting the above value for σ \sigma back to the inequality and using theorem’s assumption gives us following inequality: L 𝒟 ​ ( 𝒘 ) \displaystyle L_{\mathscr{D}}(\bm{w}) ≤ ( 1 − 1 / n ) ​ max ‖ ϵ ‖ 2 ≤ ρ ​ L 𝒮 ​ ( 𝒘 + ϵ ) + 1 / n \displaystyle\leq(1-1/\sqrt{n})\max_{\|\bm{\epsilon}\|_{2}\leq\rho}L_{\mathcal{S}}(\bm{w}+\bm{\epsilon})+1/\sqrt{n} + 1 4 ​ k ​ log ⁡ ( 1 + ‖ 𝒘 ‖ 2 2 ρ 2 ​ ( 1 + log ⁡ ( n ) k ) 2 ) + log ⁡ n δ + 2 ​ log ⁡ ( 6 ​ n + 3 ​ k ) n − 1 \displaystyle+\sqrt{\frac{\frac{1}{4}k\log\left(1+\frac{\|\bm{w}\|_{2}^{2}}{\rho^{2}}\left(1+\sqrt{\frac{\log(n)}{k}}\right)^{2}\right)+\log\frac{n}{\delta}+2\log\left(6n+3k\right)}{n-1}} ≤ max ‖ ϵ ‖ 2 ≤ ρ ⁡ L 𝒮 ​ ( 𝒘 + ϵ ) + \displaystyle\leq\max_{\|\bm{\epsilon}\|_{2}\leq\rho}L_{\mathcal{S}}(\bm{w}+\bm{\epsilon})+ + k ​ log ⁡ ( 1 + ‖ 𝒘 ‖ 2 2 ρ 2 ​ ( 1 + log ⁡ ( n ) k ) 2 ) + 4 ​ log ⁡ n δ + 8 ​ log ⁡ ( 6 ​ n + 3 ​ k ) n − 1 \displaystyle+\sqrt{\frac{k\log\left(1+\frac{\|\bm{w}\|_{2}^{2}}{\rho^{2}}\left(1+\sqrt{\frac{\log(n)}{k}}\right)^{2}\right)+4\log\frac{n}{\delta}+8\log\left(6n+3k\right)}{n-1}}

∎

## Appendix B Additional Experimental Results

### B.1 SVHN and Fashion-MNIST

We report in table 5 results obtained on SVHN and Fashion-MNIST datasets. On these datasets, SAM allows a simple WideResNet to reach or push state-of-the-art accuracy (0.99% error rate for SVHN, 3.59% for Fashion-MNIST).

For SVHN, we used all the available data (73257 digits for training set + 531131 additional samples). For auto-augment, we use the best policy found on this dataset as described in ( Cubuk et al., 2018 ) plus cutout ( Devries & Taylor, 2017 ) . For Fashion-MNIST, the auto-augmentation line correspond to cutout only.

## Appendix C Experiment Details

### C.1 Hyperparameters for Experiments

We report in table 6 the hyper-parameters selected by gridsearch for the CIFAR experiments, and the ones for SVHN and Fashion-MNIST in 7 . For CIFAR-10, CIFAR-100, SVHN and Fashion-MNIST, we use a batch size of 256 and determine the learning rate and weight decay used to train each model via a joint grid search prior to applying SAM; all other model hyperparameter values are identical to those used in prior work.

For the Imagenet results (ResNet models), the models are trained for 100, 200 or 400 epochs on Google Cloud TPUv3 32 cores with a batch size of 4096. The initial learning rate is set to 1.0 and decayed using a cosine schedule. Weight decay is set to 0.0001 with SGD optimizer and momentum = 0.9.

Finally, for the noisy label experiments, we also found ρ \rho by gridsearch, computing the accuracy on a (non-noisy) validation set composed of a random subset of 10% of the usual CIFAR training samples. We report the validation accuracy of the bootstrapped version of SAM for different levels of noise and different ρ \rho in table 8 .

### C.2 Finetuning Details

Weights are initialized to the values provided by the publicly available checkpoints, except the last dense layer, which change size to accomodate the new number of classes, that is randomly initialized. We train all models with weight decay 1 ​ e − 5 1e^{-5} as suggested in ( Tan & Le, 2019 ) , but we reduce the learning rate to 0.016 as the models tend to diverge for higher values. We use a batch size of 1024 on Google Cloud TPUv3 64 cores and cosine learning rate decay. Because other works train with batch size of 256, we train for 5k steps instead of 20k. We freeze the batch norm statistics and use them for normalization, effectively using the batch norm as we would at test time 11 11 11 We found anecdotal evidence that this makes the finetuning more robust to overtraining. . We train the models using SGD with momentum 0.9 and cosine learning rate decay. For Efficientnet-L2, we use this time a batch size 512 to save memory and adjusted the number of training steps accordingly. For CIFAR, we use the same autoaugment policy as in the previous experiments. We do not use data augmentation for the other datasets, applying the same preprocessing as for the Imagenet experiments. We also scale down the learning rate to 0.008 as the batch size is now twice as small. We used Google Cloud TPUv3 128 cores. All other parameters stay the same. For Imagenet, we trained both models from checkpoint for 10 epochs using a learning rate of 0.1 and ρ = 0.05 \rho=0.05 . We do not randomly initialize the last layer as we did for the other datasets, but instead use the weights included in the checkpoint.

### C.3 Experimental results with ρ = 0.05 \rho=0.05

A big sensitivity to the choice of hyper-parameters would make a method less easy to use. To demonstrate that SAM performs even when ρ \rho is not finely tuned, we compiled the table for the CIFAR and the finetuning experiments using ρ = 0.05 \rho=0.05 . Please note that we already used ρ = 0.05 \rho=0.05 for all Imagenet experiments. We report those scores in table 9 and 10 .

### C.4 Ablation of the Second Order Terms

As described in section 2, computing the gradient of the sharpness aware objective yield some second order terms that are more expensive to compute. To analyze this ablation more in depth, we trained a WideResNet-40x2 on CIFAR-10 using SAM with and without discarding the second order terms during training. We report the cosine similarity of the two updates in figure 5 , along the training trajectory of both experiments. We also report the training error rate (evaluated at 𝒘 + ϵ ^ ​ ( 𝒘 ) \bm{w}+\hat{\bm{\epsilon}}(\bm{w}) ) and the test error rate (evaluated at 𝒘 \bm{w} ).

We observe that during the first half of the training, discarding the second order terms does not impact the general direction of the training, as the cosine similarity between the first and second order updates are very close to 1. However, when the model nears convergence, the similarity between both types of updates becomes weaker. Fortunately, the model trained without the second order terms reaches a lower test error, showing that the most efficient method is also the one providing the best generalization on this example. The reason for this is quite unclear and should be analyzed in follow up work.

### C.5 Choice of p-norm