Reweighting Local Mimina with Tilted SAM

Authors: Tian Li, Tianyi Zhou, Jeff Bilmes
Primary area: optimization
Keywords: sharpness-aware optimization, exponential tilting, generalization

Abstract:
Sharpness-Aware Minimization (SAM) has been demonstrated to improve the generalization performance of overparameterized models by seeking flat minima on the loss landscape through optimizing model parameters that incur the largest loss within a neighborhood. 
Nevertheless, such min-max formulations are computationally challenging especially when the problem is highly non-convex. Additionally, focusing only on the worst-case local solution while ignoring potentially many other local solutions may be suboptimal when searching for flat minima. In this work, we propose Tilted SAM (TSAM), a generalization of SAM inspired by exponential tilting that effectively assigns higher priority to local solutions that are flatter and that incur larger losses. TSAM is parameterized by a tilt hyperparameter $t$ and reduces to SAM as $t$ approaches infinity. We prove that (1) the TSAM objective is smoother than SAM and thus easier to optimize; and (2) TSAM explicitly favors flatter minima as $t$ increases. This is desirable as flatter minima could have better generalization properties for certain tasks. We develop algorithms motivated by the discretization of Hamiltonian dynamics to solve TSAM. Empirically, TSAM arrives at  flatter local minima and results in superior test performance than the baselines of SAM and ERM across a range of image and text tasks.

Decision: reject
Reviewer scores: 6, 6, 8, 5

Reviews and discussion:
<thread>
creation date:
1729460887497

thread_type:
official_review

<thread_object>
creator: reviewer_agf9

<content>
summary:
The paper presents Tilted SAM (TSAM) that generalizes Sharpness-Aware Minimization (SAM) by incorporating exponential tilting. While SAM aims to minimize the worst-case loss in a neighborhood around the model parameters to attain flatter minima, TSAM enhances this approach by emphasizing local solutions that are flatter but allow for larger losses, controlled by a tilt hyperparameter $t$. Compared to SAM, TSAM is less aggressive as it adjusts the perturbation between worst-case scenarios and other states. The paper includes theoretical proof of TSAM's properties, describes an algorithm inspired by Hamiltonian dynamics for solving TSAM, and presents empirical results demonstrating TSAM's superior generalization performance and ability to achieve flatter minima.

strengths:
1. TSAM, a smoother objective for SAM, achieves flatter minima through the use of exponential tilting.

2. The paper defines sharpness in terms of loss variance and provides proof that TSAM can achieve flatter minima.

3. The authors propose a practical algorithm inspired by Hamiltonian Monte Carlo.

4. Experimental results show that TSAM achieves better generalization and attains flatter minima in both image and language classification tasks.

weaknesses:
1. Section 3.2 demonstrates that TSAM prefers a flatter model as $t$ increases. However, the relationship between flatness and generalization is still unclear as the authors mentioned in related works. Although the authors analyze the generalization of TSAM in Section 3.3, it's a bit weak since Remark 1 just claims that TSAM _can_ obtain a smaller linear population error of ERM or SAM.

2. Lack of experiments on large-scale datasets like ImageNet-1K, which is a standard dataset in SAM-related works.

3. In Figure 3, it seems $t$ is a sensitive hyperparameter and it needs to be tuned on different datasets, which adds extra computational cost. In addition, as $t$ increases that should lead to flatter minima as discussed in Section 3.2, but it doesn't always improve the generaliation. How to pick effective $t$ becomes a problem.

4. Why choose $\mathrm{E}[L(\theta^\star + \epsilon)]$ and $\mathrm{var}[L(\theta^\star + \epsilon)]$ representing sharpness in Section 5.2? Hessian spectra, the maximum eigenvalue, and $\lambda_1/\lambda_5$ are the most common methods to describe the sharpness in SAM-related works (e.g. SAM paper and VASSO paper, both of them mentioned in this paper). Therefore, the authors should explain the reason, and that would be better if adding additional results.

5. Computational cost is too much. TSAM requires 3 or 5 times the computational cost of SAM, which makes it impractical. Don't even mention the costs of tuning $t$ and $s$.

6. It's not very clear about experimental settings. For example, what's the noise level of Noisy CIFAR100? Are the results of Table 1 mean values, how many runs of each experiment, and could you provide the standard deviation?

questions:
1. What's the meaning of '(1)' at the end of line 281?

2. What's the role of $\hat{\theta}$ in Algorithm 2? Why not compute $\epsilon$ directly?

3. Is it a typo in line 475? When $t=0$, TSAM reduces to SAM objective or average weight perturbation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730672082858

thread_type:
official_review

<thread_object>
creator: reviewer_uh13

<content>
summary:
This paper presents an approach named Tilted Sharpness-Aware Minimization (TSAM), which extends the Sharpness-Aware Minimization (SAM) method to improve generalization in overparameterized machine learning models. SAM optimizes model parameters by focusing on minimizing the maximum loss in a neighborhood, aiming to find flat minima that enhance generalization. However, SAM's approach can be computationally challenging and overly focused on the worst-case local solutions. 

TSAM addresses these issues by introducing a "tilt" parameter formulation $t$, which allows for reweighting local minima, focusing on flatter solutions that could generalize better. The authors demonstrate that as the tilt parameter increases, TSAM smooths the optimization landscape, favoring flatter minima. To effectively solve the TSAM objective, the authors propose a novel Hamiltonian Monte Carlo-inspired algorithm to sample and estimate gradients efficiently. Empirical evaluations across image and text tasks show TSAM’s superiority in test performance over SAM and ERM, including fine-tuning ResNet18 and Vision Transformers on CIFAR-100, ImageNet, and DTD datasets, and fine-tuning DistilBert model on the GLUE benchmark

strengths:
- TSAM’s approach to generalization via tilting provides a novel and flexible framework, allowing it to optimize for both worst-case and average-case scenarios, unlike SAM’s rigid min-max formulation.Theoretical proofs and empirical results support that TSAM achieves flatter minima than SAM, correlating with improved generalization on test data.

- TSAM is tested on a diverse range of tasks and models (e.g., CNNs, transformers, and GLUE benchmark tasks), demonstrating its robustness and adaptability to different domains.

- By using a Hamiltonian Monte Carlo method to sample neighborhood perturbations, TSAM optimizes for flatter minima more efficiently than naive sampling would allow, enabling practical deployment in high-dimensional optimization tasks.

weaknesses:
- The algorithm introduces additional parameters (e.g., tilt parameter $t$ and sampling frequency), requiring careful tuning, which could be challenging in resource-constrained environments.

- Although TSAM is more efficient than a straightforward min-max approach, it still requires multiple gradient evaluations per iteration. Can the authors analyze the runtime complexity of the optimization algorithm? Moreover, it would be better to provide a comparison of running time between the baselines in the paper.

questions:
How can one tune the hyper-parameter $t$? It would be better to provide a clear method/intuition for selecting an optimal $t$.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730710943425

thread_type:
official_review

<thread_object>
creator: reviewer_wkwd

<content>
summary:
The paper introduces a new version of the sharpness-aware minimization algorithm (SAM) called tilted SAM (TSAM). 
TSAM uses exponential tilting to smooth the objective. TSAM comes with a temperature parameter t.
 As t approaches infinity TSAM converges to classic SAM.
The paper shows exponential tilting leads to an easier optimization problem while still favoring flat solutions.
TSAM's theoretical analysis shows that it enjoys good generalization properties.
Moreover, the numerical experiments show TSAM performs superior to other SAM-style algorithms on several benchmark tasks.

strengths:
**Clarity**

The paper is well-presented and easy to read.

**Theory**

The paper provides good theoretical guarantees for the algorithm, establishing that the solution found by the algorithm generalizes well under standard hypotheses.

**Experiments**

TSAM performs better than existing SAM-like approaches on standard ML tasks. 
Moreover, the competitors are tuned, which is good. 

**Overall** 

This is a solid paper and could be of interest to certain subcommunities working on deep learning theory and optimization.

weaknesses:
**Technical Contribution**

The argument for establishing the main result, Theorem 2, seems standard. It consists of a standard application of Hoeffding's inequality combined with a technical lemma proved by Aminian et al. (2024).

**Numerical Experiments**

The paper only reports performance over epochs and not wall clock time. It should report both; if the proposed algorithm runs significantly slower than existing methods, that's problematic.

For the HMC portion of the algorithm, to generate $\epsilon$, the authors mention they set $N=1$ for computational efficiency.
Given the experimental results, this approximation seems to work. 
Nevertheless, it would have been nice to have some ablation in the appendix with larger values of $N$ to show the impact increasing this parameter has on performance.
This would help validate the setting of $N=1$ and show if there is any tradeoff between speed and performance in setting $N$.

questions:
1. Could you add wall clock times in the revision?
2. Could you add an ablation to the appendix for the $N$ parameter.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730729037202

thread_type:
official_review

<thread_object>
creator: reviewer_1ja6

<content>
summary:
Inspired by the idea of exponential tilting in probability and statistics, this paper introduces a generalized and smoothed variant of SAM, termed TSAM. The authors highlight two key attributes of TSAM: smoothness and preference for flatter minima. They further develop an algorithm to practically solve TSAM. Finally, they verify the generalization improvement of TSAM over SAM and ERM via extensive experiments.

strengths:
1. The idea of reweighting multiple points in the neighborhoods of the local minima instead of focusing on the single worst solution is intuitive.

2. The paper is well-written and supports its claims through theory and comprehensive experiments.

weaknesses:
1. There seem to be conflicting explanations for why TSAM outperforms SAM regarding generalization. On one side, the authors suggest that TSAM improves generalization by considering a broader range of local solutions rather than concentrating solely on the worst-case scenarios, indicating that TSAM's objective function is inherently superior for identifying solutions with better generalization. Conversely, TSAM is said to favor flatter minima as $t$ increases. As $t$ approaches infinity, TSAM converges to SAM, suggesting that SAM might ideally be the best at finding flatter minima, if not for practical optimization challenges. Thus, TSAM's enhanced performance in practice may stem from its balanced trade-off between optimizing for flatter minima and maintaining optimization efficiency.

2. The generalization bounds presented in Section 3.3 are confusing regarding the message they aim to convey. The authors seek to demonstrate that TSAM has smaller upper bounds on generalization error compared to SAM and ERM. However, in their theoretical framework, the bound for SAM is worse than that for ERM, which contradicts the superior generalization performance of SAM in practice. This contradiction undermines the validity of the comparative advantage attributed to TSAM over both SAM and ERM.

3. The improvement of TSAM over SAM and ERM is marginal in the transformer experiments (the second table in Table 1). Moreover, TSAM requires more computation than SAM and ERM, thus raising the question: if variants such as ESAM1 or ESAM2 were to be tested, would the marginal improvements provided by TSAM diminish further?

questions:
See weaknesses
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces Tilted Sharpness-Aware Minimization, a generalization of SAM that smooths the objective using exponential tilting. While the reviewers acknowledged the theoretical exploration and the intuitive motivation behind TSAM, they raised significant concerns that remain unresolved.

Key concerns include:

1) Marginal Empirical Improvement: Although TSAM shows improvements over SAM and ESAM on standard benchmarks, the gains (1-2%) are considered minor, particularly in light of the substantial computational overhead.
2) Computational Cost: TSAM requires 3-5 times the computational time and memory of SAM, making it impractical for large-scale datasets.
3) Theoretical Contributions: While the analysis supports TSAM’s preference for flatter minima, reviewers noted that the bounds rely on standard results and do not provide new insights into the relationship between flatness and generalization.
4) Hyperparameter Sensitivity: The tilt parameter requires careful tuning, further increasing computational costs.

The authors provided clarifications during the rebuttal, but the reviewers felt the issues regarding practical utility, computational efficiency, and theoretical significance were not fully resolved. Consequently, I recommend rejection.