Reweighting Local Mimina with Tilted SAM

We thank the reviewer for the comments and questions.

[Generalization] Theorem 2 implies that TSAM (under some bounded $t$) would result in a smaller upper bound of linear population error compared with ERM, which we think is an interesting result. As we compare the upper bounds of the linear population risk (which is the metric of interest in practical deployment), optimizing for TSAM is optimizing for a smaller upper bound. We provide results (Theorem 1) where TSAM leads to better generalization or flatness; but the relationship between flatness and generalization in deep neural works is still an open problem (Wen et al., 2023).

[Wen et al., 2023] Sharpness Minimization Algorithms Do Not Only Minimize Sharpness To Achieve Better Generalization, NeurIPS 2023.

[Additional experiments] Thanks for the suggestion of evaluating TSAM on an additional ImageNet dataset. Due to constraints with the setup of our compute resources, we find it’s going to take a long time to run imagenet experiments with proper hyperparameter tuning for all methods. Therefore, we explore the tiny imagenet dataset, which is a subset of imagenet with 200 classes. In the table below, we report the results on finetuning tiny imageNet with ResNet18 starting from a pre-trained model. For TSAM, we sample $s$=3 perturbations. Note that ESAM1 is simply the vanilla SAM update but running $3\times$ more iterations. We see that TSAM again outperforms the baselines on this dataset. This trend is consistent with other results on a variety of image and language benchmarks and models in the paper.

[How to tune t] In some cases (DTD and noisy CIFAR100 datasets), there exist a wide range of $t$’s that all result in improvements compared with SAM. But there is an optimal $t$ that leads to the best test performance, which needs manual tuning. In our experiments, we only pick $t$ from a limited set of grids of {0, 1, 5, 20, 100}. Though tuning $t$ requires extra effort, we obtain better generalization on a variety of tasks. Note that for the baselines (e.g., more advanced SAM variants such as PGN or RSAM) that TSAM outperforms, they generally introduce additional one or two new hyperparameters that we also tuned carefully. One goal of this work is to provide some high-level intuitions around the effects of $t$ (Section 3).

[Sharpness definition] Thanks for the suggestion of evaluating more sharpness definitions. As partially mentioned in Section 5.2, we choose mean and variance because (a) those are intuitive statistics to reason about flatness of local loss surface, (b) they have appeared in prior works called average-perturbation sharpness or average-direction sharpness (e.g., Chen et al., 2022, Wen et al., 2023), and (c) such notions are more aligned with our theoretical analysis in Section 3.2.

Having said that, we agree that it would be better to evaluate on more sharpness measures. We have computed the eigenvalues of the hessian under SAM and TSAM solutions (fixing the same computation budget) of the CIFAR100 dataset trained ResNet18. Top-5 eigenvalues are {342.11, 304.72, 260.71, 252.92, 210.88} for ERM, {232.60, 198.35, 182.61, 153.74, 145.76} for SAM, and {140.91, 113.38, 105.90, 92.94, 89.55} for TSAM ($t$=20). We see that TSAM achieves the smallest max eigenvalues among the three objectives. We have added more discussions around this and the results to the revision.

[Chen et al., 2022] When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations, ICLR 2022. [Wen et al., 2023] Sharpness Minimization Algorithms Do Not Only Minimize Sharpness To Achieve Better Generalization, NeurIPS 2023.

We thank the reviewer for the time and effort spent and letting us know the remaining concerns. We have revised the paper around these comments, and our responses are provided below.

[Presentation] Thanks for the suggestion. Part of the confusion may come from the fact that understanding the relations between flatness and generalization is still an active area of research on its own. In our work, we do not claim theoretically flatness leads to generalization in our context. Rather, we proved that TSAM helps with both flatness and generalization to the extent described in Theorem 1 and Theorem 2. We have revised Section 5.2 and moved Table 2 to the appendix in the revision.

[Additional experiments] We note that while we did not use several model architectures for each dataset, we already have a broad coverage of models (CNNs, vision transformers, and bert models), datasets (noisy and clean versions, three image and one text benchmarks), and evaluation (various baselines, training and test performance, and sharpness of the final solutions) in experiments. We also added tiny imagenet results in Table 4 in the Appendix. Due to computation resource constraints, we will provide complete imagenet results in the next version.

We appreciate the reviewer’s additional comment of evaluating multiple architectures for each dataset. Hence, we perform extra evaluation on the WideResNet model on three datasets for all baselines. Test accuracies are shown in the table below, and have also been added to the revision of the paper (Table 1, Section 5.1).

We see that TSAM (number of perturbations $s$=3) still demonstrates superior performance than the baselines under the same computational budget (e.g., TSAM running for 200 epochs and ESAM1 running for 600 epochs). If we continue to increase the runtime of the baselines, their performance wouldn’t improve.

As discussed in our paper (end of the first paragraph of Section 2), as TSAM is a new sharpness-aware formulation, it’s possible to apply many optimization techniques to solve TSAM, e.g., combining the current TSAM algorithms with other SAM variants. One way is to apply adaptivity or variance reduction on top of the tilted gradients.

We have also included the papers the reviewers mentioned in the paper’s references.

[Computational cost of TSAM] We agree that in order to estimate the tilted gradients for TSAM, we generally need more gradient evaluations per iteration, as discussed and acknowledged in the paper. However, we have evaluated TSAM and all the baselines with the same computational budget, e.g., by letting baselines run 3$\times$ or 5$\times$ longer than TSAM, or optimizing for the inner max of SAM in a more fine-grained way using more gradient computation (ESAM2). We see that TSAM still outperforms those approaches on a variety of tasks and models. The trend still holds even if we further increase the runtime of ERM (SGD) or SAM. In addition, the implementation of TSAM can be optimized. If we sample the perturbations (which are independent to each other) in parallel, we can incur the same per-iteration runtime as the vanilla SAM algorithm.

We note that other SAM baselines can introduce new hyperparameters as well (e.g., the coefficient parameter in PGN, the noise scaling factor and the ratio between noise and gradients lambda in RSAM). In TSAM, we tune $t$ from a limited set of candidates, and there are multiple $t$’s that can lead to improved performance. In addition, we provide theoretical analysis on the effects of t, which we hope can provide additional intuition. As our responses to Reviewer uH13, one interesting direction is to automatically optimize $t$ as an optimization variable by taking inspirations from previous works (e.g., Eq. (65) in Li et al. (2023)) and connecting the TSAM objective with the optimization of quantile perturbed losses. We leave this for future work.

[Li et al. 2023] On Tilted Losses in Machine Learning: Theory and Applications, JMLR 2023.

Thanks for the detailed reply. I increase my score to 5.

Please ensure that experiments involving training on ImageNet-1k are included in your next version. Thanks!

We are grateful for the insightful reviews, and the acknowledgement of the theoretical and empirical contributions of our work.

[Runtime comparisons] We note that on the image datasets of Table 1, we compare the performance of various algorithms under the same computational budget. Though TSAM requires more gradient computation per iteration, it still outperforms baselines with fewer total iterations. We report the exact runtime numbers for TSAM and baselines in Table 6 of the revision. In terms of optimization complexity, if we assume we can get an unbiased estimator of the tilted gradients at each iteration, the convergence speed of the optimization algorithm is $O(1/T)$ for strongly convex and smooth functions, following standard arguments in optimization. TSAM is also easily parallelizable by sampling perturbations at the same time; we will make such implementation optimization in the next version.

[Tuning t] While we usually need to tune $t$ manually, one goal of this work is to provide some high-level intuitions around effects of $t$. For instance, we discuss smoothness, convexity, flatness, and generalization as functions of t in Section 3. It is an interesting future direction to explore how to automatically optimize $t$ in an end-to-end way. One high-level idea is to take inspirations from previous works (e.g., Eq. (65) in Li et al. (2023)) and connect the TSAM objective with the optimization of quantile perturbed losses. It is thus possible to reformulate the problem, taking $t$ as an optimization variable as opposed to a hyperparameter. We leave this for future work. Thanks for bringing this up!

[Li et al., 2023] On Tilted Losses in Machine Learning: Theory and Applications, JMLR 2023.

We thank the reviewer for the positive review of our work.

[Wall clock times] Though we did not report the exact runtime of the experiments, we compare different methods under the same computation time in Table 1 (our main result), including expensive versions of SAM (ESAM1 and ESAM2) and letting other sharpness-aware variants (PGN and RSAM) run longer than TSAM. We show that under the same runtime, TSAM outperforms other algorithms (expensive SAM and other SAM variants). We have added the exact wall clock time in the revision (Table 6, Appendix 4.2) and highlighted them in blue.

[Effects of N] Our empirical results in the original submission are based on running $N$=1 steps in HMC (Algorithm 1 and Algorithm 2). Per the reviewer’s suggestion, we still report the performance with $N$>1 in Table 4, Appendix C.2 of the revision. We assume we accept all the generated epsilons for simplicity and easier comparisons. For different $N$’s, we sample the same number of perturbations. The result is also copied in the table below. We observe diminishing returns when increasing $N$.

[Others] Thanks for the reviewer’s comments on “...this paper provides good theoretical guarantees…”. Regarding the proof of Theorem 2, we agree that the proof is not technically complicated, which we view as a strength, since we arrive at meaningful results with clean proofs built upon previous works. Theorem 2 implies that there exists an optimal $t$ that minimizes the upper bound of the generalization error defined as linear population risk. Additionally, we provide various major theoretical arguments in Section 3 (e.g., smoothness, convexity, and flatness of solutions) in addition to the theorem discussed here to help understand the properties of TSAM.

​​

[Flatness and generalization of TSAM solutions] We agree with the reviewer that TSAM outperforms SAM due to the combined effects of encouraging flat minima and the smoothness of the objective. The relations between flatness and generalization remains an interesting open problem (e.g., Wen et al., 2023). In Theorem 1, we comment on the flatness property of TSAM for generalization linear models, which by itself does not directly prove that TSAM generalizes better than SAM. Regarding the generalization bound in Theorem 2, the upper bound of $E_Z[l(\theta, Z)]$ (the linear population risk that we are concerned about) converges to the classic generalization result of ERM as $t$ converges to 0. Hence, our result is relatively tight for a bounded $t$, implying the superiority of TSAM over ERM. We agree that the bound in its current form does not explain the superiority of SAM over ERM, as the optimal $t$ that minimizes the upper bound is not $\infty$. We have revised the corresponding paragraph in the revision to avoid confusion.

[Wen et al., 2023] Sharpness Minimization Algorithms Do Not Only Minimize Sharpness To Achieve Better Generalization, NeurIPS 2023.

[Empirical improvements] Thanks for the suggestion of evaluating other baselines on the language tasks. We note that compared with numbers reported in previous works (e.g., Bahri et al., Table 2), our improvements are on a similar scale (if we convert the results in our paper to the form of percentage). We observe that on bert experiments, finetuning can be done typically within a small number of epochs (e.g., 5 to 10 depending on the dataset in the Glue benchmark). If we continue to finetune using SAM (ESAM1), the performance will not improve over SAM. For the ESAM1 and ESAM2 baselines, the results on the Glue benchmark are as follows. TSAM still demonstrates higher scores. We have also added these results in the second table of Table 1.

[Bahri et al., 2022] Sharpness-Aware Minimization Improves Language Model Generalization, ACL 2022.

Thank you for clarifying the advantages of TSAM over SAM and for doing the additional experiments on ESAM. These new experiments addressed my concerns regarding the effectiveness of TSAM. However, as noted in Weakness 2 and by other reviewers, TSAM incurs more computational overhead compared to SAM. Given this limitation, I have decided not to raise my score, but I still lean towards acceptance.