Tilted Sharpness-Aware Minimization

We thank the reviewer for the time and valuable comments. We hope our response below can address the reviewer's concerns.

[more experiment results] We appreciate the reviewer’s suggestions for adding more baselines. We have cited the papers mentioned by the reviewer in related work. We did not compare them in our original submission mainly because they are proposed to solve a SAM-like minimax style objective, where TSAM proposes a new objective function than SAM along with an algorithm to solve it. As plotted in Figures 3 and 4 in the appendix, even if we can solve the minimax SAM objective perfectly (in that particular toy problem, we did a brute-force search of the optimal parameters in the one-dimensional space), TSAM solutions still demonstrate more desirable properties than SAM solutions (more smooth).

However, we agree more results will strengthen our experiments, and we compare with both GAM [3] and GSAM [2]. We leave out ASAM because GSAM achieves better performance than it (Figure 5 of the GSAM paper). For GSAM [2], we tune all the hyperparameters including the alpha parameter from a grid of {0.01, 0.02, 0.03} as suggested by the paper. The GAM algorithm [3] can be expensive as it requires to compute Hessian vector products. We implemented its non-expensive approximation, as explained in the GAM paper and actually implemented in its open-sourced code.

Results are shown in the table below. We report the best performance of each baseline on each dataset. For instance, in noisy CIFAR100, both GSAM and GAM overfit to noise at the later stage of training and we apply early stopping to obtain the best test performance.

[1] ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks, ICML 2021.

[2] Surrogate Gap Minimization Improves Sharpness-Aware Training, ICLR 2022.

[3] Gradient Norm Aware Minimization Seeks First-Order Flatness and Improves Generalization, CVPR 2023.

[compute cost] We agree that TSAM requires 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 (Section 5), 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, TSAM is inherently easy to parallelize since all the perturbations of model parameters are independent of each other. If we optimize our implementation by sampling perturbations in parallel, TSAM would incur the same per-iteration runtime as the vanilla SAM algorithm.

[typos] Thanks for pointing out the typos in the submission. We will fix them in the next version.

We thank the reviewer for the time and valuable comments.

[expensive computation] We would like to clarify that the runtime numbers in Table 5 are reported for the same number of iterations of all methods (to illustrate the per-iteration cost), while the final test performance in the experiment section (Section 5) is obtained by letting all methods run the same amount of wall-clock time. Across all datasets, we consistently observe that under the same runtime, TSAM outperforms other algorithms.

In addition, TSAM is inherently easy to parallelize since all the perturbations of model parameters are independent of each other. If we optimize our practical implementation in this way, TSAM would incur the same per-iteration runtime as the vanilla SAM algorithm (consisting of one-step of gradient ascent and one step of gradient descent each iteration).

[scheduling $t$] We note that scheduling $t$ with a fixed constraint set (defining the neighborhood of model parameters) does not change final performance significantly as long as we start from or end with the same values of $t$. When $t$ is zero, TSAM reduces to an average-perturbation-based sharpness-aware objective, which still accounts for some notions of sharpness of the loss surface. In the future, we plan to further investigate scheduling $t$ together with the radius parameter of the constraint set where it is possible to cover vanilla ERM with a mini-batch SGD optimizer.

We thank the reviewer for the time and positive assessment of our work! We would like to address the remaining questions/concerns as follows.

[experimental improvements] We observe statistically significant improvements of TSAM compared with the baselines. In terms of the test loss, TSAM outperforms baselines by a large margin (Table 3 in Appendix C.3). Furthermore, as TSAM is a new objective function minimizing a weighted combination of bad losses, we can further improve the current algorithm by combining it with other optimization techniques such as applying the idea of variance reduction to obtain better estimated gradients, or incorporate adaptivity to precondition the estimated gradients. We leave more fine-grained exploration of solvers as future work.

[other measures of sharpness and related work] Thanks for the suggestions and for pointing out the related work. In our work, we have explored two metrics of sharpnesses, as discussed in Section 5.2. The first sharpness measure (results visualized in Table 2) is closely related to the local entropy notion mentioned in the related work (Pittorino et. al.) and we agree the formulations are conceptually related as well. We will add discussions around the connections with Pittorino et. al. Our formulation is inspired by exponential tilting, which inherently has rich connections with different areas such as information theory, applied probability, and optimization [1]. We are certainly interested in developing other sharpness measures based on the properties of our objectives leveraging connections with prior works, and we will add discussions on this in the next version.

[1] On Tilted Losses in Machine Learning: Theory and Applications, JMLR.

We thank the reviewer for the time and valuable feedback.

[connections between flat local minima and improved generalization] Note that we do not intend to claim ‘flat local minima leads to a better generalization in all cases’ by claiming “TSAM arrives at flatter local minima and results in superior test performance…” The exact relations between the two for deep learning is still an open question, as pointed out in our submission citing several related works (Section 2). Our study is partly motivated by the empirical success of a series of SAM works under which sharpness of local minima is significantly reduced (under various notions of sharpness), and we aim to develop a better formulation along this line by considering and reweighting multiple local minima in a principled framework. We will further clarify this in the next version.

[runtime of TSAM] We agree that TSAM requires 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 (Section 5), 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 complicated way with 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, TSAM is inherently easy to parallelize since all the perturbations of model parameters are independent of each other. If we optimize our implementation in this way, TSAM would incur the same per-iteration runtime as the vanilla SAM algorithm.

[other questions] (a) The performance of both TSAM and SAM would improve very slightly after 200 epochs, but TSAM still outperforms SAM. (b) When $N=1$, we are running gradient ascent once to locate the regions with (relatively) large losses, where the first-order gradient information serves as a strong signal to guide the areas to sample from. Additionally, the gap between losses can be magnified by using a large value of $t$ during reweighting. We will explain the intuition in more detail in the next version.