Tilted Sharpness-Aware Minimization

TL;DR

This paper introduces Tilted SAM (TSAM), a new optimization method that generalizes SAM by prioritizing flatter minima, leading to better generalization in overparameterized models across image and text tasks.

Contribution

TSAM is a novel, smoothed variant of SAM that explicitly favors flatter minima and is easier to optimize, improving generalization performance.

Findings

TSAM finds flatter minima than SAM and ERM.

TSAM achieves superior test performance on image and text tasks.

TSAM is smoother and easier to optimize than SAM.

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 smoothed generalization of SAM inspired by exponential tilting that effectively assigns higher priority to local solutions that incur larger losses. TSAM is parameterized by a tilt hyperparameter $t$ and reduces to SAM as $t$ approaches infinity. We show that TSAM is smoother than SAM and thus easier to optimize, and it explicitly favors flatter minima. 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.