# A Universal Class of Sharpness-Aware Minimization Algorithms

Behrooz Tahmasebi
MIT CSAIL
<EMAIL_ADDRESS>Ashkan Soleymani
MIT LIDS
<EMAIL_ADDRESS>Dara Bahri
Google DeepMind
<EMAIL_ADDRESS>Stefanie Jegelka
TU Munich and MIT CSAIL
<EMAIL_ADDRESS>Patrick Jaillet
MIT LIDS
<EMAIL_ADDRESS>
###### Abstract

Recently, there has been a surge in interest in developing optimization
algorithms for overparameterized models as achieving generalization is
believed to require algorithms with suitable biases. This interest centers on
minimizing sharpness of the original loss function; the Sharpness-Aware
Minimization (SAM) algorithm has proven effective. However, most literature
only considers a few sharpness measures, such as the maximum eigenvalue or
trace of the training loss Hessian, which may not yield meaningful insights
for non-convex optimization scenarios like neural networks. Additionally, many
sharpness measures are sensitive to parameter invariances in neural networks,
magnifying significantly under rescaling parameters. Motivated by these
challenges, we introduce a new class of sharpness measures in this paper,
leading to new sharpness-aware objective functions. We prove that these
measures are universally expressive, allowing any function of the training
loss Hessian matrix to be represented by appropriate hyperparameters.
Furthermore, we show that the proposed objective functions explicitly bias
towards minimizing their corresponding sharpness measures, and how they allow
meaningful applications to models with parameter invariances (such as scale-
invariances). Finally, as instances of our proposed general framework, we
present Frob-SAM and Det-SAM, which are specifically designed to minimize the
Frobenius norm and the determinant of the Hessian of the training loss,
respectively. We also demonstrate the advantages of our general framework
through extensive experiments.

##  1 Introduction

Understanding the generalization capabilities of overparameterized networks is
a fundamental, yet unsolved challenge in deep learning. It is postulated that
achieving near-zero training loss alone may be insufficient as there exist
many instances where global minima fail to exhibit satisfactory generalization
performance. To this end, a dominant observation asserts that the
characteristics of the loss landscape play a pivotal role in determining which
parameters have low training loss while also exhibiting generalization
capabilities.

A recently proposed approach to consider the geometric aspects of the loss
landscape, with the aim of achieving generalization, entails the avoidance of
sharp minima. For example, the celebrated Sharpness-Aware Minimization (SAM)
algorithm has shown enhancements in generalization across many practical tasks
(Foret et al., 2021). While the concept of sharpness lacks a precise
definition in a general sense, people often introduce various measures to
quantify it in practice (Dinh et al., 2017). Many sharpness measures in the
literature rely on the second-order derivative characteristics of the training
loss function, such as the trace or the operator norm of the Hessian matrix
(Chaudhari et al., 2019; Keskar et al., 2017).

Nevertheless, traditional methodologies for quantifying sharpness may not
suffice to ensure generalization, given the intricate geometry of the loss
landscape, which may necessitate different regularization techniques.
Moreover, many existing sharpness measures fail to encapsulate the genuine
essence of sharpness in deep neural networks because the Hessian matrix no
longer maintains positive semi-definiteness. Furthermore, neural networks
exhibit parameter invariances, wherein different parameterizations can yield
identical functions — such as scaling invariances in ReLU networks.
Consequently, an effective measure of sharpness should remain invariant in the
face of such parameter variations. Unfortunately, conventional approaches for
quantifying sharpness frequently fall short in addressing this phenomenon.

Therefore, a fundamental question arises: how can one succinctly represent all
measures of sharpness within a compact parameterized framework that also
enables meaningful applications to models with parameter invariances? This
question holds significance in applications as it allows learning/designing
the regularization in cases where information about the geometry of the loss
landscape or parameter invariances is provided, either empirically or through
assumption. To the best of our knowledge, this question has remained fairly
unexplored in the deep learning literature.

In this paper, we characterize all sharpness measures (i.e., functions of the
Hessian of the training loss) through an average-based parameterized
representation. We prove that by changing the (hyper)parameters, the provided
representation spans all the sharpness measures as a function of the Hessian
matrix. In other words, it is provably a universal representation. We also
provide quantitative theoretical results on the complexity of the sharpness
representation as a function of the data dimension.

Moreover, attached to any representation of sharpness, we provide a new loss
function, and we prove that the new loss function is biased toward minimizing
its corresponding sharpness measure. Since the parameterized representation
reduces to SAM (i.e., worst-direction) and average-direction sharpness
measures (Wen et al., 2023a) in special cases, it can be considered as a
generalized (hyper)parameterized sharpness-aware minimization algorithm. This
generalizes the recent study of the explicit bias of a few sharpness-aware
minimization algorithms (Wen et al., 2023a) to a comprehensive class of
objectives. Furthermore, this allows us to readily design algorithms with any
bias of interest, while to the best of our knowledge, only algorithms with
biases towards minimizing the trace, operator norm of the Hessian matrix, and
a few other sharpness measures are known in the literature. As instances of
our proposed general algorithm, we present Frob-SAM and Det-SAM, two new
sharpness-aware minimization algorithms that are specifically designed to
minimize the Frobenius norm and the determinant of the Hessian of the training
loss function, respectively.

An interesting feature of the given representation is that it provides a
systematic way to construct sharpness measures respecting parameter
invariances, e.g., scale-invariance in neural networks. As a specific example,
we provide a class of loss functions and algorithms that are invariant under
parameter scaling. Note that the classical sharpness measures, such as trace
or the operator norm of the Hessian matrix, are not invariant to rescaling or
group actions.

In the experiments, we explore extensively two specific choices of these
algorithms: (1) Frob-SAM: an algorithm biased toward minimizing the Frobenius
norm of the Hessian, a meaningful sharpness notion for non-convex optimization
problems and (2) Det-SAM: an algorithm biased toward minimizing the
determinant of the Hessian, a scale-invariant sharpness measure. We
demonstrate the advantages of these two cases through an extensive series of
experiments.

In short, in this paper we make the following contributions:

  * $\bullet$

We propose a new class of sharpness measures, as function of the training loss
Hessian. We prove that the new representation is universally expressive,
meaning that it covers all sharpness measures of the Hessian as its
(hyper)parameters change.

  * $\bullet$

Along with each sharpness measure we provide an optimization objective and
prove that the new objective is explicitly biased toward minimizing the
corresponding sharpness measure.

  * $\bullet$

The structure of the proposed method allows meaningful applications to models
with parameter invariances, as it provides a class of objective functions for
any desired type of parameter invariance.

  * $\bullet$

We introduce two fundamental illustrative examples of our proposed general
representation and the corresponding algorithms: Frob-SAM and Det-SAM. Frob-
SAM is geared towards minimizing the Frobenius norm of the Hessian matrix,
providing a meaningful and natural solution to the definition problem of
sharpness for non-convex optimization problems. Conversely, Det-SAM is focused
on minimizing the determinant111To be more precise, the product of non-zero
eigenvalues. of Hessian, addressing scale-invariant issues related to
parameterization.

##  2 Related Work

Foret et al. (2021) recently proposed the Sharpness-Aware Minimization (SAM)
algorithm to avoid sharp minima. The SAM objective has connections to a
similar robust optimization problem that was suggested for the study of
adversarial attacks in deep learning (Madry et al., 2018). Besides SAM,
Nitanda et al. (2023) show how parameter averaging for SGD is biased toward
flatter minima. Label noise SGD also prefers flat minima (Damian et al.,
2021). Woodworth et al. (2020) studied the role of sharpness in
overparameterization from a kernel perspective. See (Wang et al., 2023) for
the applications of flat minima for domain generalization (see also (Cha et
al., 2021)). For applications of SAM in large language models, see (Bahri et
al., 2022) (also (Zhong et al., 2022), and (Qu et al., 2022; Shi et al., 2023)
for federated learning). Besides those applications, Wen et al. (2023b) prove
that current sharpness minimization algorithms sometimes fail to generalize
for non-generalizing flattest models.

The (implicit) bias of many optimization algorithms and architectures has been
studied, from the Gradient Descent (GD) (Ji and Telgarsky, 2019b; Soudry et
al., 2018) to mirror descent (Gunasekar et al., 2018a; Azizan and Hassibi,
2019); see also (Gunasekar et al., 2018b) for linear convolutional networks,
and (Lawrence et al., 2022) for equivariant networks. Ji and Telgarsky (2019a)
observed that linear neural networks are biased toward weight alignment for
different layers (see (Le and Jegelka, 2022) for non-linear networks).
Andriushchenko and Flammarion (2022) study implicit bias of SAM for diagonal
linear networks, and Wen et al. (2023a) find the explicit bias of the Gaussian
averaging method and other SAM variants.

The role of scale-invariance in generalization in deep learning is emphasized
in (Neyshabur et al., 2017). Dinh et al. (2017) point out that parameter
invariances can lead to the different parameterization of the same function,
making the definition of flatness challenging; see also (Andriushchenko et
al., 2023b) for a recent study. This motivates the study of sharpness measures
that are invariant to such reparametrizations.

There have been a few attempts to address reparametrization problems with
sharpness measures recently. Kwon et al. (2021) proposed to adaptively
calculate the sharpness in a normalized ball around the loss function to
achieve scale invariance. However, their method is limited to scaling
problems. Kim et al. (2022) took a step further and introduced a new SAM
algorithm by capturing the neighborhood of the parameters in an ellipsoid
induced by the Fisher information. This way, the neighborhood becomes
invariant with respect to the parameter invariances in the network. Jang et
al. (2022) defined an information geometric sharpness measure by investigating
the eigenspaces of Fisher Information Matrix (FIM) of distribution
parameterized by neural networks. They proved scale-invariance properties for
their notion. Even though Kim et al. (2022) and Jang et al. (2022) enjoy some
parameter invariance properties, (1) in practice, their methods are limited to
classification tasks because of FIM calculation, (2) the underlying explicit
biasing of their algorithms remains a mystery and is not guaranteed.222Please
refer to Appendix A for a more detailed overview of related work.

##  3 Background

###  3.1 Setting

Consider a standard learning setup with a labeled dataset $\mathcal{S}$, and a
training loss function $L:\mathbb{R}^{d}\to\mathbb{R}_{\geq 0}$, where $L(x)$
denotes the training loss over $\mathcal{S}$ computed for the parameters
$x\in\mathbb{R}^{d}$. The main objective in Empirical Risk Minimization (ERM)
is to minimize the training loss $L(x)$ over the feasibility set
$\mathcal{X}\subseteq\mathbb{R}^{d}$. However, achieving parameters satisfying
$L(x)\approx 0$ in overparameterized models is often straightforward. This is
because in contrast to other models, in overparameterized models, there are
many global minima, i.e., the set $\Gamma:=\\{x\in\mathcal{X}:L(x)=0\\}$ is a
manifold – it is called the zero-loss manifold in the literature. Moreover, in
practical scenarios, it is noteworthy that not all global minima exhibit
favorable generalization capabilities (Foret et al., 2021).

###  3.2 Background on SAM

It is hypothesized that the avoidance of sharp minima can help generalization
performance (Hochreiter and Schmidhuber, 1997; Keskar et al., 2017; Izmailov
et al., 2018). However, it should be noted that the concept of sharpness
encompasses a multitude of distinct definitions in practical contexts. The
Sharpness-Aware Minimization (SAM) algorithm (Foret et al., 2021) suggests
minimizing the training loss function over a small ball around the parameters:

$\displaystyle\min_{x\in\mathcal{X}}\Big{\\{}L_{\text{SAM}}(x):=\max_{\|v\|_{2}\leq
1}L(x+\rho v)\Big{\\}},$

where $\rho\in\mathbb{R}_{\geq 0}$ is the _perturbation parameter_. Note that
$L_{\text{SAM}}$ can be decomposed into two terms:

$\displaystyle L_{\text{SAM}}(x)=\underbrace{L(x)}_{\text{empirical
loss}}+\underbrace{\max_{\|v\|_{2}\leq 1}\big{\\{}L(x+\rho
v)-L(x)\big{\\}}}_{\text{sharpness}}.$

Foret et al. (2021) also suggest alternative average-based sharpness-aware
objectives to use PAC bounds on the generalization error of overparameterized
models; we follow the definition in (Wen et al., 2023b):

$\displaystyle L_{\text{AVG}}(x)$
$\displaystyle:=\mathbb{E}_{v\sim\mathcal{N}(0,I)}\Big{[}L(x+\frac{\rho
v}{\|v\|_{2}})\Big{]}=\underbrace{L(x)}_{\text{empirical
loss}}+\underbrace{\mathbb{E}_{v\sim\mathcal{N}(0,I)}\Big{[}L(x+\frac{\rho
v}{\|v\|_{2}})-L(x)\Big{]}}_{\text{sharpness}}.$

Wen et al. (2023b) recently proved that minimizing $L_{\text{SAM}}(x)$ will
lead to global minima (i.e., $L(x)\approx 0$) with small
$\lambda_{\max}(\nabla^{2}L(x))$. In other words, SAM is (explicitly) biased
towards minimizing $\lambda_{\max}(\nabla^{2}L(x))$. Moreover, they show that
using $L_{\text{AVG}}(x)$ biases towards minimizing
$\frac{1}{d}\operatorname*{tr}(\nabla^{2}L(x))$. This means that SAM measures
the sharpness of a global minimum by $\lambda_{\max}(\nabla^{2}L(x))$, while
the average-based objective uses
$\frac{1}{d}\operatorname*{tr}(\nabla^{2}L(x))$ to evaluate it.

In the next examples, we argue how both sharpness measures above fail to
define a meaningful notion for overparameterized models. In Example 2, a
special case of problem with parameter invariances, i.e., under parameter
rescalings is discussed.

######  Example 1.

The sharpness measures
$\lambda_{\max}(\nabla^{2}L(x)),\lambda_{\min}(\nabla^{2}L(x))$ and
$\operatorname*{tr}(\nabla^{2}L(x))$ are conceptually meaningful when the
objective function $L(x)$ is convex, therefore $\lambda_{i}$s are nonnegative.
However, the Loss landscape of neural networks is highly nonconvex, and as a
result, $\lambda_{i}$ can be potentially negative. Consider the toy non-convex
example of $L(x_{1},x_{2})=\frac{1}{2}(x_{1}^{2}-x_{2}^{2}$),

$\displaystyle\nabla^{2}L=\begin{bmatrix}1&0\\\ 0&-1\end{bmatrix}.$ (1)

For all $x_{1},x_{2}\in\mathbb{R}$, we know that
$\textup{tr}(\nabla^{2}L)=1+(-1)=0$, which in the Trace measure of sharpness
it suggests that all the points $x_{1},x_{2}\in\mathbb{R}$ are equally flat.
Are these sharpness notions really capturing the intended concepts? For a
better illustration, consider the plot of this function provided in Figure 1.
This problem extends to other existing notions.

Figure 1: The loss landscape of the non-convex objective function
$L(x_{1},x_{2})=\frac{1}{2}(x_{1}^{2}-x_{2}^{2})$. This examples shows how
existing sharpness measures fall short of capturing sharpness meaning in non-
convex settings. In particular, for all points
$(x_{1},x_{2})\in\mathbb{R}^{2}$,
$\textup{tr}(\nabla^{2}L(x_{1},x_{2}))=1+(-1)=0$.

######  Example 2.

Consider the loss function $L(x_{1},x_{2})=x_{1}^{2}x_{2}^{2}-2x_{1}x_{2}+1$
with two parameters $x_{1},x_{2}\in\mathbb{R}$. It is scale-invariant, i.e.,
$L(kx_{1},\frac{x_{2}}{k})=L(x_{1},x_{2})$ for all $k\neq 0$. Indeed, the
zero-loss manifold $\Gamma=\\{(x_{1},x_{2}):x_{1}x_{2}=1\\}$ contains
infinitely many global minima. Straightforward calculation shows
$\nabla^{2}L(x_{1},x_{2})=\big{(}\begin{smallmatrix}2x_{2}^{2}&4x_{1}x_{2}-2\\\
4x_{1}x_{2}-2&2x_{1}^{2}\end{smallmatrix}\big{)}$. Thus, we have
$\frac{1}{2}\operatorname*{tr}(\nabla^{2}L(x_{1},x_{2}))=x_{1}^{2}+x_{2}^{2}$.
After rescaling, we get
$\frac{1}{2}\operatorname*{tr}(\nabla^{2}L(x_{1},x_{2}))\Big{|}_{(kx_{1},k^{-1}x_{2})}=k^{2}x_{1}^{2}+\frac{x_{2}^{2}}{k^{2}}\neq\frac{1}{2}\operatorname*{tr}(\nabla^{2}L(x_{1},x_{2}))$.
Therefore, as a sharpness measure,
$\operatorname*{tr}(\nabla^{2}L(x_{1},x_{2}))$ is not scale-invariant. The
problem magnifies in the limit:
$\operatorname*{tr}(\nabla^{2}L(x_{1},x_{2}))\Big{|}_{(kx_{1},k^{-1}x_{2})}\to\infty$
as $k\to\infty$. Similar problems exist for
$\lambda_{\max}(\nabla^{2}L(x_{1},x_{2}))$. However,
$\det(\nabla^{2}L(x_{1},x_{2}))$ is scale-invariant; we have
$\det(\nabla^{2}L(x_{1},x_{2}))\Big{|}_{(kx_{1},k^{-1}x_{2})}=\det(\nabla^{2}L(x_{1},x_{2}))$
for all $k\neq 0$.

Note that neural networks are often scale-invariant, e.g., linear networks or
ReLU networks after scaling up the parameters of one hidden layer and scaling
down the parameters of another hidden layer encode the same functions.

##  4 A New Class of Sharpness Measures

To define a new class of sharpness measures, we take a closer look at the
average-based sharpness-aware objective $L_{\text{AVG}}(x)$; using its Taylor
expansion (Wen et al., 2023b), we have

$\displaystyle L_{\text{AVG}}(x)$
$\displaystyle=\mathbb{E}_{v\sim\mathcal{N}(0,I)}\Big{[}L(x+\frac{\rho
v}{\|v\|_{2}})\Big{]}$ $\displaystyle\approx
L(x)+\rho\mathbb{E}_{v\sim\mathcal{N}(0,I)}\Big{[}\langle\nabla
L(x),\frac{v}{\|v\|_{2}}\rangle\Big{]}+\rho^{2}\mathbb{E}_{v\sim\mathcal{N}(0,I)}\Big{[}\frac{v^{t}\nabla^{2}L(x)v}{\|v\|^{2}_{2}}\Big{]}$
$\displaystyle=L(x)+\rho^{2}\frac{\operatorname*{tr}(\nabla^{2}L(x))}{d}.$

This intuitively tells us that for a small perturbation parameter $\rho$, the
leading term in the objective function is the training loss $L(x)$, and after
we get close to the zero-loss manifold $\Gamma$, the leading term becomes
$\frac{1}{d}\operatorname*{tr}(\nabla^{2}L(x))$, which is exactly the explicit
bias of the average-based sharpness-aware minimization objective. This
motivates us to define the following parameterized sharpness measure.

######  Definition 1 ($(\phi,\psi,\mu)$-sharpness measure).

For any continuous functions $\phi,\psi:\mathbb{R}\to\mathbb{R}$ and any
(Borel) measure $\mu$ on $\mathbb{R}^{d}$, the $(\phi,\psi,\mu)$-sharpness
measure $S(x;\phi,\psi,\mu)$ is defined as

$\displaystyle
S(x;\phi,\psi,\mu):=\phi\Big{(}\int\psi\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)\Big{)}.$
(2)

Similarly, one can consider continuous functions
$\psi:\mathbb{R}\to\mathbb{R}^{m}$ and $\phi:\mathbb{R}^{m}\to\mathbb{R}$, for
some positive integer $m\geq 1$, and (Borel) measures $\mu_{\ell}$,
$\ell\in[m]$, and define

$\displaystyle S(x;\phi,\psi,\mu):=\phi\Big{(}$
$\displaystyle\int\psi_{1}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu_{1}(v),$
(3)
$\displaystyle\int\psi_{2}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu_{2}(v),$
(4) $\displaystyle\ldots,$ (5)
$\displaystyle\int\psi_{m}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu_{m}(v)\Big{)},$
(6)

where we use $\mu:=\mu_{1}\otimes\mu_{2}\otimes\ldots\otimes\mu_{m}$ for the
sake of brevity in our notation, and
$\psi=(\psi_{1},\psi_{2},\ldots,\psi_{m})^{t}$.

We specify several examples of hyperparameters $(\phi,\psi,\mu)$ in Table 1,
which shows how $(\phi,\psi,\mu)$-sharpness measures can represent various
notions of sharpness, as a function of the training loss Hessian matrix.

Table 1:  Various $(\phi,\psi,\mu)$-sharpness measures (Appendix B;
$\lambda_{i}$, $i\in[d]$, are eigenvalues of Hessian).

Hyperparameters |  |
---|---|---
$\phi(t)$ | $\psi(t)$ | $m$ | $\mu$ |  $S(x;\phi,\psi,\mu)$ (or bias) | Reference
$t$ | $t$ | $1$ | $\text{Uniform}(\mathbb{S}^{d-1})$ | $\frac{1}{2d}\operatorname*{tr}(\nabla^{2}L(x))=\frac{1}{2d}\sum_{i=1}^{d}\lambda_{i}$ | (Wen et al., 2023a)
$2\max(\\{.\\})$ | $(t,t,\ldots)$ | $\infty$ | $\otimes_{\|v\|_{2}=1}\delta_{v}$ | $\max_{i\in[d]}\lambda_{i}$ | (Wen et al., 2023a)
$2(t_{2}-t_{1}^{2})$ | $(t,t^{2})$ | $2$ | $\mathcal{N}(0,I_{d})\otimes\mathcal{N}(0,I_{d})$ | $\sum_{i=1}^{d}\lambda^{2}_{i}=\|\nabla^{2}L(x)\|_{F}$ | This paper
$(2\pi)^{d}/t^{2}$ | $\exp(-t)$ | $1$ | Lebesgue measure on $\mathbb{R}^{d}$ | $\det(\nabla^{2}L(x))=\prod_{i=1}^{d}\lambda_{i}$ | This paper
$t$ | $t^{n}$ | $1$ | $\text{Uniform}(\mathbb{S}^{d-1})$ | $q(\lambda_{1},\lambda_{2},\ldots,\lambda_{d})$ | $\text{This paper}^{*}$
$1/t^{2}$ | $\exp(\sigma t)$ | $1$ | $\mathcal{N}(0,I_{d})$ | $\prod_{i=1}^{d}(1-\sigma\lambda_{i})$ | This paper

*$q(\lambda_{1},\lambda_{2}\ldots,\lambda_{d})$ is a specific homogeneous polynomial of degree $n$; see Equation 22.

##  5 Expressive Power and Universality

In this section, we prove that the proposed class of sharpness measures is
universal. In other words, for any continuous function
$S:\mathbb{R}^{d}\to\mathbb{R}$, we specify continuous functions $\phi,\psi$
and a (Borel) probability measure333We indeed prove that Borel probability
measures (as a subset of arbitrary Borel measures) are enough to achieve
universality.  $\mu$ on $\mathbb{R}^{d}$ such that
$S(\lambda_{1},\lambda_{2}\ldots,\lambda_{d})=S(x;\phi,\psi,\mu)$, where
$\lambda_{i}$, $i\in[d]$, are the eigenvalues of the Hessian matrix
$\nabla^{2}L(x)$.

######  Theorem 1 (Universality of the $(\phi,\psi,\mu)$-sharpness measures
for functions of Hessian eigenvalues).

Let $\mathcal{A}\subseteq\mathbb{R}^{d}$ be a compact set. For any continuous
function $S:\mathcal{A}\to\mathbb{R}$, there exist a product (Borel)
probability measure $\mu$, a positive integer $m\leq d$, and continuous
functions $\phi:\mathbb{R}^{m}\to\mathbb{R}$ and
$\psi:\mathbb{R}\to\mathbb{R}^{m}$, such that
$S(\lambda_{1},\lambda_{2}\ldots,\lambda_{d})=S(x;\phi,\psi,\mu)$ for any
$x\in\mathcal{A}$, where $\lambda_{i}$, $i\in[d]$, are the eigenvalues of the
Hessian matrix $\nabla^{2}L(x)$.

We present the proof of Theorem 1 in Appendix C.

Note that to achieve universality, we need the functions $\phi,\psi$ to be of
dimension $m=d$. However, as one can see in Table 1, many celebrated sharpness
measures can indeed be represented using only small $m$. We believe that
practically small hyperparameter $m$ is enough, as it is motivated from the
measures in Table 1.

While we proved the universality of the proposed class of sharpness measures
for continuous functions of the Hessian eigenvalues, one may be interested in
measuring sharpness with more information about the loss Hessian (e.g., the
eigenvectors of the loss Hessian). The following theorem proves the
universality for this class of arbitrary functions.

######  Theorem 2 (Universality of the $(\phi,\psi,\mu)$-sharpness measures
for arbitrary functions of Hessian).

For any continuous function $S:\mathbb{R}^{d\times d}\to\mathbb{R}$, there
exist a positive integer $m\leq d(d+1)/2$, (Borel) probability measures
$\mu_{\ell}$, $\ell\in[m]$, and continuous functions
$\phi:\mathbb{R}^{m}\to\mathbb{R}$ and $\psi:\mathbb{R}\to\mathbb{R}^{m}$,
such that $S(\nabla^{2}L(x))=S(x;\phi,\psi,\mu)$ for any $x\in\mathbb{R}^{d}$,
where $\mu:=\mu_{1}\otimes\mu_{2}\otimes\ldots\otimes\mu_{m}$ is a product
probability measure.

We present the proof of Theorem 2 in Appendix D.

Note that arbitrary functions of the Hessian matrix can be quite hard to
compute, e.g., consider the permanent of the Hessian matrix. Moreover, the
dimension $m$ must be quite large to allow us to prove the universality in
overparameterized models (for $d$ of considerable size), since the generality
bound scales as $\mathcal{O}(d^{2})$. Nevertheless, in practice, only small
$m$ allows to cover many interesting cases.

##  6 Explicit Bias

Now that we defined a flexible set of sharpness measures and we proved that it
is universally expressive, the following question arises: how can one achieve
$S(x;\phi,\psi,\mu)$ as the explicit bias of an objective function that only
relies on the zeroth-order information about the training loss, similar to
$L_{\text{SAM}}(x)$ and $L_{\text{AVG}}(x)$? To answer this question, we
introduce the $(\phi,\psi,\mu)$-sharpness-aware loss function as follows.

######  Definition 2.

The $(\phi,\psi,\mu)$-sharpness-aware loss function

$\displaystyle L_{(\phi,\psi,\mu)}$
$\displaystyle(x):=\underbrace{L(x)}_{\text{empirical
loss}}+\rho^{2}\underbrace{\phi\Big{(}\int\psi\Big{(}\frac{1}{\rho^{2}}\big{(}L(x+\rho
v)-L(x)\big{)}\Big{)}d\mu(v)\Big{)}}_{:=R_{\rho}(x)\leavevmode\nobreak\
\text{sharpness}}=L(x)+\rho^{2}R_{\rho}(x),$

where $\rho$ is the perturbation parameter and $R_{\rho}(x)$ denotes the
sharpness regularizer.

Extending this definition to the cases with $m>1$ is straightforward.

In the above definition, the new regularizer $R_{\rho}(x)$ is an approximation
of the sharpness measure $S(x;\phi,\psi,\mu)$ as $\rho\to 0^{+}$. As a result,
it is expected that minimizing $L_{(\phi,\psi,\mu)}(x)$ lead to minimizing the
training loss as well as the sharpness measure $S(x;\phi,\psi,\mu)$. The next
theorem formalizes this intuitive observation via characterizing the explicit
bias of minimizing the sharpness-aware loss function $L_{(\phi,\psi,\mu)}(x)$.

######  Theorem 3 (Explicit bias of the $(\phi,\psi,\mu)$-sharpness-aware loss
function).

Given a triplet $(\phi,\psi,\mu)$, $m\geq 1$, and a training loss function
$L:\mathbb{R}^{d}\to\mathbb{R}_{\geq 0}$, assume that:

  * •

$L(x)$ is third-order continuously differentiable and satisfies the following
upper bound

$\displaystyle\max_{i,j,k\in\\{1,2,3\\}}|\partial_{i}\partial_{j}\partial_{k}L(v)|=O(\|v\|^{-1}),$
(7)

for $v\in\mathbb{R}^{d}$ as $\|v\|_{2}\to\infty$.

  * •

The two functions $\phi,\psi$ are continuously differentiable.

  * •

For some
$C>\max_{x\in\mathcal{X}}\max_{i\in[d]}|\lambda_{i}(\nabla^{2}L(x))|$, we have
$\int\|v\|_{2}^{2}\tilde{\psi}_{i}(v)d\mu(v)<\infty$, $i\in[m]$, where

$\displaystyle\tilde{\psi}_{i}(v):=\max_{|t|\leq\|v\|_{2}}|\psi_{i}^{\prime}(Ct^{2})|.$
(8)

Then, there exists an open neighborhood $U\supsetneq\Gamma$, where $\Gamma$ is
the zero-loss manifold, for connected $U$ and $\Gamma$, such that if for some
$u\in U$, one has

$\displaystyle L(u)+\rho^{2}R_{\rho}(u)-\inf_{x\in
U}\Big{(}L(x)+\rho^{2}R_{\rho}(x)\Big{)}\leq\Delta\rho^{2},$ (9)

with some optimally gap $\Delta>0$, then

$\displaystyle L(u)\leq\underbrace{\inf_{x\in
U}L(x)}_{=0}+(\Delta+o_{\rho}(1))\rho^{2},$ (10)

and also

$\displaystyle
S(u;\phi,\psi,\mu)\leq\inf_{x\in\Gamma}S(x;\phi,\psi,\mu)+\Delta+o_{\rho}(1).$
(11)

We present the proof of Theorem 3 in Appendix E.

The above theorem shows how using the new objective function
$L_{(\phi,\psi,\mu)}(x)$ leads to explicitly biased optimization algorithms
towards minimizing the sharpness measure $S(x;\phi,\psi,\mu)$ over the zero-
loss manifold $\Gamma$. Indeed, it proves that if we are close to the zero-
loss manifold (i.e., $u\in U$ for some open neighborhood $U\supsetneq\Gamma$),
and also $L_{(\phi,\psi,\mu)}(u)$ is close to its global minimum over $U$,
then (1) the training loss function $L(u)$ is close to zero, and (2) the
corresponding sharpness measure $S(u;\phi,\psi,\mu)$ is close to its global
minimum over the zero-loss manifold, with respect to an optimality gap
$\Delta$.

Algorithm 1 $(\phi,\psi,\mu)$-Sharpness-Aware Minimization Algorithm (with
$m=1$)

Input: The triplet $(\phi,\psi,\mu)$, training loss $L(x)$, step size $\eta$,
perturbation parameter $\rho$, number of samples $n$,

Output: Model parameters $x_{t}$ trained with $(\phi,\psi,\mu)$-sharpness-
aware minimization algorithm

Initialization: $x\leftarrow x_{0}$ and $t\leftarrow 0$

while 1 do

Sample $v_{1},v_{2},\ldots,v_{n}\overset{\text{i.i.d.}}{\sim}\mu$

Compute the following:  $\displaystyle g_{t}$ $\displaystyle=\nabla
L(x_{t})+\phi^{\prime}\Big{(}\sum_{i=1}^{n}\frac{1}{n}\psi\Big{(}\frac{1}{\rho^{2}}\big{(}L(x_{t}+\rho
v_{i})-L(x_{t})\big{)}\Big{)}\Big{)}$ $\displaystyle\hskip
56.9055pt\times\sum_{i=1}^{n}\frac{1}{n}\Big{\\{}\psi^{\prime}\Big{(}\frac{1}{\rho^{2}}\big{(}L(x_{t}+\rho
v_{i})-L(x_{t})\big{)}\Big{)}\times\Big{(}\nabla L(x_{t}+\rho v_{i})-\nabla
L(x_{t})\Big{)}\Big{\\}}.$

Update the parameters: $x_{t+1}=x_{t}-\eta g_{t}$

$t\leftarrow t+1$

end while

##  7 Invariant Sharpness-Aware Minimization

For which hyperparameters $(\phi,\psi,\mu)$ is the corresponding sharpness
measure scale-invariant? The following theorem answers this question.

######  Theorem 4 (Scale-invariant $(\phi,\psi,\mu)$-sharpness measures).

Consider a scale-invariant loss function $L(x)$ and let $\mu$ be a Borel
measure of the form

$\displaystyle
d\mu(x)=f\Big{(}\prod_{i=1}^{d}x_{i}\Big{)}\prod_{i=1}^{d}dx_{i},$ (12)

where $f:\mathbb{R}\to\mathbb{R}$ is a measurable function444In Lemma 1, we
show that any scale-invariant measure is of this form.. Then, for any
continuous functions $\phi,\psi$, the corresponding sharpness measure
$S(x;\phi,\psi,\mu)$ is scale-invariant; this means that
$S(x;\phi,\psi,\mu)=S(Dx;\phi,\psi,\mu)$ for any diagonal matrix
$D\in\mathbb{R}^{d\times d}$ with $\det(D)=1$.

We present the proof of Theorem 4 in Appendix F.

######  Example 3.

Note that $\det(\nabla^{2}L(x))$ is a scale-invariant sharpness measure; for
any diagonal matrix $D\in\mathbb{R}^{d\times d}$ with $\det(D)=1$,

$\displaystyle\det(\nabla^{2}L(x))$
$\displaystyle\Big{|}_{Dx}=\det(D^{-1}\nabla^{2}L(x)D^{-1})=\det(D^{-1})^{2}\det(\nabla^{2}L(x))=\det(\nabla^{2}L(x)).$

Note that Theorem 4 also supports the scale-invariance of the determinant; the
Lebesgue measure satisfies the condition in Theorem 4 with $f\equiv 1$, and we
have the representation of the determinant in Table 1.

While in Theorem 4 we only considered scale-invariances, one can generalize it
to a general class of parameter invariances in the following theorem.

######  Theorem 5 (General parameter-invariant $(\phi,\psi,\mu)$-sharpness
measures).

Let $G$ be a group acting by matrices on $\mathbb{R}^{d}$, and assume that
$L(x)$ is invariant with respect to the action of $G$. Then, for any
$G$-invariant (Borel) measure $\mu$, and any continuous functions $\phi,\psi$,
the corresponding sharpness measure $S(x;\phi,\psi,\mu)$ is $G$-invariant;
this means that $S(x;\phi,\psi,\mu)=S(A_{g}x;\phi,\psi,\mu)$ for any matrix
$A_{g}\in\mathbb{R}^{d\times d}$ corresponding to the action of an element
$g\in G$.

The proof of this theorem is analogous to Theorem 4 and is deferred to
Appendix G. Thus, the strategy to create sharpness measures invariant to any
group action $G$ is simply to choose a group action invariant measure $\mu$.
Now, for any family of choices of functions $\phi$ and $\psi$, we obtain a
family of $G$-invariant sharpness measures. Consequently, there is a family of
$G$-invariant Sharpness-Aware Minimization algorithms, as explained in Section
8.

##  8 $(\phi,\psi,\mu)$-Sharpness-Aware Minimization Algorithm

In this section, we present the pseudocode for the
$(\phi,\psi,\mu)$-Sharpness-Aware Minimization Algorithm (see Algorithm 1).
For simplicity, we present the algorithm for the full-batch gradient descent,
and assume that $m=1$. Extending it to the mini-batch case with $m>1$ is
straightforward (see Algorithm 3). The idea is to apply (stochastic) gradient
decent or other optimization algorithms on the $(\phi,\psi,\mu)$-sharpness-
aware loss function defined in Definition 2,

$\displaystyle L_{(\phi,\psi,\mu)}=L(x)+\rho^{2}R_{\rho}(x).$

However, calculating the sharpness term $R_{\rho}(x)$ directly is analytically
hard to do because of the integration with respect to the probability measure
$\mu$. Hence, we propose to estimate the inner integration at each iteration
with i.i.d. random variables $\nu_{1},\nu_{2},\dots,\nu_{n}\sim\mu$ as
perturbations, i.e.,

$\displaystyle\tilde{R}_{\rho}(x)\coloneqq\phi\Big{(}\frac{1}{n}\sum_{i=1}^{n}\psi\Big{(}\frac{1}{\rho^{2}}\big{(}L(x+\rho
v_{i})-L(x)\big{)}\Big{)}\Big{)}.$

When $\phi$ satisfies continuity conditions, for large enough $n$, the
estimator $\tilde{R}_{\rho}(x)$ will converge to $R_{\rho}(x)$. Now, we
calculate the gradients of $L(x)+\rho^{2}\tilde{R}_{\rho}(x)$. By chain rule,

$\displaystyle\rho^{2}\nabla\tilde{R}_{\rho}(x)=\phi^{\prime}\Big{(}\sum_{i=1}^{n}\frac{1}{n}\psi\Big{(}\frac{1}{\rho^{2}}\big{(}L(x_{t}+\rho
v_{i})-L(x_{t})\big{)}\Big{)}\Big{)}$
$\displaystyle\times\sum_{i=1}^{n}\frac{1}{n}\Big{\\{}\psi^{\prime}\Big{(}\frac{1}{\rho^{2}}\big{(}L(x_{t}+\rho
v_{i})-L(x_{t})\big{)}\Big{)}$ $\displaystyle\times\Big{(}\nabla L(x_{t}+\rho
v_{i})-\nabla L(x_{t})\Big{)}\Big{\\}},$

which leads to Algorithm 1.

Our algorithm needs $n+1$ gradient evaluations per iteration, which for $n=1$
matches the SAM algorithm (Foret et al., 2021). In practice, small values for
$n$ demonstrate the expected results, therefore, the computational overhead of
our algorithm is not a barrier.

Note that to recover the original SAM algorithm, one can set the function
$\phi,\psi$ to identity, $m=1$, and choose $\mu$ to be the single-point
measure on $\nabla L(x_{t})/\|\nabla L(x_{t})\|_{2}$ with $n=1$ sample for
each $t$.

Moreover, even though to prove universality, we only used probability
measures, we proposed a compact representation of determinant with Lebesgue
measure with $m=1$ in Table 1 and Theorem 4. However, integrals with respect
to Lebesgue measure cannot be estimated via sampling and we need to truncate
the integral to integration over a large hypercube; this allows us to use
Algorithm 1 for the scale-invariant sharpness measures. Also, this
approximation achieves non-zero sharpness in cases that the Hessian matrix is
not full-rank (which happens in overparametrized models), as it gets the
product of non-zero eigenvalues. We use this approximation in the next section
to implement the method.

##  9 Frobenius SAM and Determinant SAM

To be more concrete, we specify Algorithm 3 (for arbitrary $m$) to the case
with the Frobenius norm regularizes (with $m=2$), which we call the Frob-SAM
algorithm. Note that to achieve this, one needs to specify
$\phi(t_{1},t_{2})=2(t_{2}-t_{1}^{2})$ and $\psi(t)=(t,t^{2}).$ Furthermore,
since we only need to collect samples from the Gaussian distribution to get
the Frobenius norm bias (see Table 1), we can use the same samples to estimate
both integrals for the functions $\psi_{1}(t)=t$ and $\psi_{2}(t)=t^{2}$.
Replacing these assumptions into the formula given in Algorithm 3, we get the
following update rule:

$\displaystyle g_{t}=\nabla L(x_{t})+$ $\displaystyle
4\sum_{i=1}^{n}\frac{1}{n\rho^{2}}\Big{\\{}\big{(}L(x_{t}+\rho
v_{i})-L(x_{t})\big{)}\times\Big{(}\nabla L(x_{t}+\rho v_{i})-\nabla
L(x_{t})\Big{)}\Big{\\}}$
$\displaystyle-4\Big{\\{}\sum_{i=1}^{n}\frac{1}{n\rho}\big{(}L(x_{t}+\rho
v_{i})-L(x_{t})\big{)}\Big{\\}}\times\Big{\\{}\sum_{j=1}^{n}\frac{1}{n\rho}\Big{(}\nabla
L(x_{t}+\rho v_{j})-\nabla L(x_{t})\Big{)}\Big{\\}}.$

If we take a closer look at this, we observe that

$\displaystyle g_{t}=\nabla
L(x_{t})+\frac{4}{\rho^{2}}\widehat{\text{cov}}\Big{(}\big{(}L(x_{t}+\rho
v)-L(x_{t})\big{)},\big{(}\nabla L(x_{t}+\rho v)-\nabla
L(x_{t})\big{)}\Big{)},$

where $\widehat{\text{cov}}$ denotes the (biased) empirical cross-covariance
between the scalar random variable $L(x_{t}+\rho v)-L(x_{t})$ and the vector-
values random variable $\nabla L(x_{t}+\rho v)-\nabla L(x_{t})$, for
$v\sim\mathcal{N}(0,I_{d})$. Since the covariance is not sensitive to the
means of random variables/vectors, we can further simply the update rule to

$\displaystyle g_{t}=\nabla
L(x_{t})+\frac{4}{\rho^{2}}\widehat{\text{cov}}\Big{(}L(x_{t}+\rho v),\nabla
L(x_{t}+\rho v)\Big{)}.$

We can further replace the unbiased estimator of the cross-covariance instead
of $\widehat{\text{cov}}$ which leads to Algorithm 2.

Algorithm 2 Frob-SAM

Input: Training loss $L(x)$, step size $\eta$, perturbation parameter $\rho$,
number of samples $n$,

Output: Model parameters $x_{t}$ trained with Frobenius SAM

Initialization: $x\leftarrow x_{0}$ and $t\leftarrow 0$

while 1 do

Sample $v_{i}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(0,I_{d})$ for any
$i\in[n]$

Compute the following:  $\displaystyle g_{t}=\nabla L(x_{t})$
$\displaystyle+4\sum_{i=1}^{n}\frac{1}{(n-1)\rho^{2}}L(x_{t}+\rho v_{i})\nabla
L(x_{t}+\rho v_{i})$
$\displaystyle-4\sum_{i=1}^{n}\frac{1}{(n-1)\rho}L(x_{t}+\rho
v_{i})\times\sum_{i=1}^{n}\frac{1}{n\rho}\nabla L(x_{t}+\rho v_{i}).$

Update the parameters: $x_{t+1}=x_{t}-\eta g_{t}$

$t\leftarrow t+1$

end while

Moreover, achieving Det-SAM is also similar to Frob-SAM, but the only
difficulty is that it involves computing an integral with respect to the
Lebesgue measure which can be challenging (Table 1). To address this issue, we
instead sample a point from the hypercube $[-t,t]^{d}$ for a hyperparameter
$t\in\mathbb{R}$ to approximate the Lebesgue measure.

##  10 Experiments

Figure 2: Sharpness measure while training on MNIST for different amounts of regularization $\lambda$. When training with Frob-SAM ($\rho=0.01$, $n=2$), the Frobenius norm of the Hessian decreases throughout training and larger $\lambda$ results in lower values of the norm. One standard error bar is shaded. All $\lambda$’s achieve over 94% final test accuracy. | CIFAR10 | CIFAR100 | SVHN | CIFAR10-S | CIFAR100-S | SVHN-S
---|---|---|---|---|---|---
Frob | $94.96^{\pm 0.05}$ | $77.16^{\pm 0.16}$ | $96.26^{\pm 0.05}$ | $\textbf{74.49}^{\pm 1.15}$ | $\textbf{38.30}^{\pm 0.89}$ | $89.23^{\pm 0.32}$
Trace | $95.04^{\pm 0.07}$ | $77.61^{\pm 0.08}$ | $96.37^{\pm 0.04}$ | $\textbf{74.46}^{\pm 0.39}$ | $\textbf{37.50}^{\pm 0.78}$ | $\textbf{89.90}^{\pm 0.32}$
Det | $95.10^{\pm 0.03}$ | $77.64^{\pm 0.12}$ | $96.32^{\pm 0.05}$ | $\textbf{74.22}^{\pm 0.80}$ | $\textbf{37.48}^{\pm 0.42}$ | $\textbf{89.85}^{\pm 0.13}$
SSAM | $\textbf{95.78}^{\pm 0.05}$ | $78.42^{\pm 0.19}$ | $\textbf{96.49}^{\pm 0.04}$ | $73.39^{\pm 0.48}$ | $35.33^{\pm 0.55}$ | $\textbf{89.87}^{\pm 0.27}$
ASAM | $95.49^{\pm 0.09}$ | $\textbf{78.91}^{\pm 0.06}$ | $96.18^{\pm 0.05}$ | $\textbf{74.80}^{\pm 0.64}$ | $\textbf{37.40}^{\pm 0.66}$ | $89.21^{\pm 0.11}$
SAM | $95.54^{\pm 0.06}$ | $78.62^{\pm 0.02}$ | $\textbf{96.45}^{\pm 0.03}$ | $72.96^{\pm 0.77}$ | $36.87^{\pm 0.47}$ | $\textbf{89.80}^{\pm 0.19}$
SGD | $94.81^{\pm 0.07}$ | $77.00^{\pm 0.11}$ | $96.06^{\pm 0.08}$ | $\textbf{74.98}^{\pm 0.92}$ | $\textbf{37.64}^{\pm 0.88}$ | $\textbf{89.89}^{\pm 0.27}$
Table 2: Final test accuracy and standard errors for the full and 10% sub-sampled datasets. Entries within two standard errors of the best are bolded. Our biases particularly help when data is limited. For example, Frob-SAM achieves over 1.5% higher test accuracy than SAM on CIFAR10-S. | CIFAR10-C | CIFAR100-C | SVHN-C
---|---|---|---
Frob | $\textbf{87.57}^{\pm 0.10}$ | $65.73^{\pm 0.22}$ | $90.11^{\pm 0.20}$
Trace | $87.20^{\pm 0.12}$ | $65.29^{\pm 0.12}$ | $\textbf{91.01}^{\pm 0.44}$
Det | $83.61^{\pm 0.06}$ | $64.73^{\pm 0.29}$ | $\textbf{90.56}^{\pm 0.09}$
SSAM | $86.24^{\pm 0.08}$ | $67.06^{\pm 0.14}$ | $\textbf{90.92}^{\pm 0.12}$
ASAM | $85.88^{\pm 0.12}$ | $\textbf{67.78}^{\pm 0.04}$ | $89.54^{\pm 0.12}$
SAM | $85.42^{\pm 0.08}$ | $66.12^{\pm 0.13}$ | $\textbf{90.66}^{\pm 0.16}$
SGD | $82.78^{\pm 0.20}$ | $64.16^{\pm 0.22}$ | $87.78^{\pm 0.16}$
Table 3: Final test accuracy and standard errors when the labels of 20% of the
training examples are corrupted. Entries within two standard errors of the
best are bolded. We see that our bias can provide a boost in this scenario.
For example, Frob-SAM achieves over 2% higher accuracy than SAM on CIFAR10-C.

The goal of our experiments is twofold. Firstly, we validate Theorem 3 by
showing that minimization of the sharpness-aware loss defined in Definition 2
and codified in Algorithm 1 has the explicit bias of minimizing the sharpness
measure. Secondly, we show our method is useful in practice by evaluating it
on benchmark vision tasks. Our code is available at
https://github.com/dbahri/universal_sam.

###  10.1 Setup

We evaluate on three vision datasets: CIFAR10, CIFAR100, and SVHN. Futhermore,
we study how our _explicit_ bias may be helpful in settings that generally
benefit from regularization – specifically, when _training data is limited_
and when _training labels are noisy_. For the former, we artificially sub-
sample each original dataset, keeping only the first 10% of training samples,
and we denote these sub-sampled datasets with “-S” (i.e. CIFAR10-S). For the
latter, we choose a random 20% of training samples to corrupt, and we corrupt
these samples by flipping their labels to a different label chosen uniformly
at random over the remaining classes. We denote these datasets by “–C” (i.e.
CIFAR10-C).

Full experimental details are deferred to Appendix I; we summarize them here.
We train ResNet18 (He et al., 2016) on the datasets using momentum-SGD and a
multi-step learning rate schedule. We run each experiment under four different
random seeds. We evaluate three sharpness measures and the following other
baselines:

Frob-SAM/Trace-SAM/Det-SAM. These correspond to instances of our algorithm
where the measure is, respectively, the Frobenius norm, trace, and determinant
of the Hessian. For Det-SAM, we set $t$, half the edge width of the
approximating hypercube, to 0.01 with little tuning.

SAM. We set $\rho$ to 0.05/0.1/0.05 for CIFAR10/CIFAR100/SVHN, following Foret
et al. (2021).

Adaptive SAM (ASAM). Kwon et al. (2021) proposes a modification of SAM that is
scale-invariant. We set $\rho$ to 0.5/1/0.5 and $\eta$ to 0.01/0.1/0.01 for
CIFAR10/CIFAR100/SVHN.

Sparse SAM (SSAM). Mi et al. (2022) speeds up and improves the performance of
SAM by only perturbing important parameters, as determined via Fisher
information and sparse dynamic training. We use SSAM-F with $\rho$ set to
0.1/0.2/0.1 for CIFAR10/CIFAR100/SVHN. 50% sparsity is used with 16 samples.

Furthermore, to demonstrate that minimization of our loss has the explicit
bias we expect, we train a simple 6-layer ReLU network with 128 hidden units
on MNIST using momentum-SGD (with momentum 0.9 and learning rate 0.001) for 20
epochs. We estimate the Frobenius norm of the Hessian (see Appendix I for
details) throughout training for Frob-SAM, under a range of regularization
strengths $\lambda$.

###  10.2 Results

MNIST results are shown in Figure 2. We see that for Frob-SAM, minimization of
the loss leads to reduction in the corresponding sharpness measure and the
reduction expectantly scales proportionally with regularization strength
$\lambda$.

Results for the main datasets are shown in Table 2 and Table 3. We find that
our method nearly always outperforms SGD and is at or sometimes above par with
SAM and alternatives, especially in the noisy label and limited data settings.
For example, for noisy label CIFAR10 (CIFAR10-C), Frob-SAM achieves nearly 5%
and 2% higher final test accuracy than SGD and SAM respectively. Our findings
suggest that the explicit biases we propose can be practically useful,
especially in the noisy label and limited data scenarios, though it is often
unclear what the _best_ bias for a particular task is, a priori.

##  11 Conclusion

In this paper, we introduce a new family of sharpness measures and demonstrate
how this parameterized representation can generate many meaningful sharpness
notions (Table 1). These measures are indeed universally expressive.
Furthermore, in Theorem 3, we illustrate how the corresponding zeroth-order
objective function for each sharpness measure is explicitly biased towards
minimizing the sharpness of the training loss. Moreover, in Theorem 4, we
prove how specific (Borel) measures can lead to scale-invariant sharpness
measures, such as the determinant of the Hessian matrix. We conclude the paper
with a series of numerical experiments showcasing the efficacy of the proposed
loss function on various practical datasets. Given the broad class of
sharpness measures we proposed, an interesting future direction is to evaluate
in practice which sharpness measure/algorithm performs the best for different
dataset. Another interesting attitude is, given the universally expressivity
of our algorithm, to meta-learn the sharpness measures on the top of the
model. This could potentially yield the identification of optimal, intricate
sharpness measures tailored to specific datasets, thereby contributing to
unraveling hidden secrets of overparameterized models.

## Acknowledgments

The authors appreciate Joshua Robinson for his insightful comments and
valuable suggestions. BT and SJ are supported by the Office of Naval Research
award N00014-20-1-2023 (MURI ML-SCOPE), NSF award CCF-2112665 (TILOS AI
Institute), NSF award 2134108, and the Alexander von Humboldt Foundation. AS
and PJ are supported by the National Research Foundation Singapore and DSO
National Laboratories under the AI Singapore Programme (AISG Award No:
AISG2-RP-2020-018).

## References

  * Agarwala and Dauphin (2023) Atish Agarwala and Yann Dauphin.  SAM operates far from home: eigenvalue regularization as a dynamical phenomenon.  In _Int. Conference on Machine Learning (ICML)_ , 2023.
  * Andriushchenko and Flammarion (2022) Maksym Andriushchenko and Nicolas Flammarion.  Towards understanding sharpness-aware minimization.  In _Int. Conference on Machine Learning (ICML)_ , 2022.
  * Andriushchenko et al. (2023a) Maksym Andriushchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion.  Sharpness-aware minimization leads to low-rank features.  _arXiv preprint arXiv:2305.16292_ , 2023a.
  * Andriushchenko et al. (2023b) Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flammarion.  A modern look at the relationship between sharpness and generalization.  In _Int. Conference on Machine Learning (ICML)_ , 2023b.
  * Arora et al. (2022) Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi.  Understanding gradient descent on the edge of stability in deep learning.  In _Int. Conference on Machine Learning (ICML)_ , 2022.
  * Azizan and Hassibi (2019) Navid Azizan and Babak Hassibi.  Stochastic gradient/mirror descent: Minimax optimality and implicit regularization.  In _Int. Conference on Learning Representations (ICLR)_ , 2019.
  * Bahri et al. (2022) Dara Bahri, Hossein Mobahi, and Yi Tay.  Sharpness-aware minimization improves language model generalization.  In _Proceedings of the Annual Meeting of the Association for Computational Linguistics_ , 2022.
  * Bartlett et al. (2022) Peter L. Bartlett, Philip M. Long, and Olivier Bousquet.  The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima.  _arXiv preprint arXiv:2210.01513_ , 2022.
  * Behdin and Mazumder (2023) Kayhan Behdin and Rahul Mazumder.  On statistical properties of sharpness-aware minimization: Provable guarantees.  _arXiv preprint arXiv:2302.11836_ , 2023.
  * Behdin et al. (2022) Kayhan Behdin, Qingquan Song, Aman Gupta, David Durfee, Ayan Acharya, Sathiya Keerthi, and Rahul Mazumder.  Improved deep neural network generalization using m-sharpness-aware minimization.  _arXiv preprint arXiv:2212.04343_ , 2022.
  * Blanc et al. (2020) Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant.  Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process.  In _Conference on Learning Theory (COLT)_ , 2020.
  * Cha et al. (2021) Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park.  Swad: Domain generalization by seeking flat minima.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2021.
  * Chaudhari et al. (2019) Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann Lecun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina.  Entropy-sgd: Biasing gradient descent into wide valleys.  _Journal of Statistical Mechanics: Theory and Experiment_ , 2019(12):124018, 2019.
  * Compagnoni et al. (2023) Enea Monzio Compagnoni, Luca Biggio, Antonio Orvieto, Frank Norbert Proske, Hans Kersting, and Aurelien Lucchi.  An SDE for modeling sam: Theory and insights.  In _Int. Conference on Machine Learning (ICML)_ , 2023.
  * Damian et al. (2021) Alex Damian, Tengyu Ma, and Jason D Lee.  Label noise SGD provably prefers flat global minimizers.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2021.
  * Dinh et al. (2017) Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio.  Sharp minima can generalize for deep nets.  In _Int. Conference on Machine Learning (ICML)_ , 2017.
  * Du et al. (2022) Jiawei Du, Daquan Zhou, Jiashi Feng, Vincent Tan, and Joey Tianyi Zhou.  Sharpness-aware training for free.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Foret et al. (2021) Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur.  Sharpness-aware minimization for efficiently improving generalization.  In _Int. Conference on Learning Representations (ICLR)_ , 2021.
  * Gunasekar et al. (2018a) Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro.  Characterizing implicit bias in terms of optimization geometry.  In _Int. Conference on Machine Learning (ICML)_ , 2018a.
  * Gunasekar et al. (2018b) Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro.  Implicit bias of gradient descent on linear convolutional networks.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2018b.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016.
  * Hochreiter and Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Flat minima.  _Neural computation_ , 9(1):1–42, 1997.
  * Izmailov et al. (2018) Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson.  Averaging weights leads to wider optima and better generalization.  In _Conference on Uncertainty in Artificial Intelligence (UAI)_ , 2018\.
  * Jang et al. (2022) Cheongjae Jang, Sungyoon Lee, Frank Park, and Yung-Kyun Noh.  A reparametrization-invariant sharpness measure based on information geometry.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Ji and Telgarsky (2019a) Ziwei Ji and Matus Telgarsky.  Gradient descent aligns the layers of deep linear networks.  In _Int. Conference on Learning Representations (ICLR)_ , 2019a.
  * Ji and Telgarsky (2019b) Ziwei Ji and Matus Telgarsky.  The implicit bias of gradient descent on nonseparable data.  In _Conference on Learning Theory (COLT)_ , 2019b.
  * Jiang et al. (2020) Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio.  Fantastic generalization measures and where to find them.  In _Int. Conference on Learning Representations (ICLR)_ , 2020.
  * Kaddour et al. (2022) Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt J Kusner.  When do flat minima optimizers work?  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Keskar et al. (2017) Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.  On large-batch training for deep learning: Generalization gap and sharp minima.  In _Int. Conference on Learning Representations (ICLR)_ , 2017.
  * Kim et al. (2023) Hoki Kim, Jinseong Park, Yujin Choi, Woojin Lee, and Jaewook Lee.  Exploring the effect of multi-step ascent in sharpness-aware minimization.  _arXiv preprint arXiv:2302.10181_ , 2023.
  * Kim et al. (2022) Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales.  Fisher sam: Information geometry and sharpness aware minimisation.  In _Int. Conference on Machine Learning (ICML)_ , 2022.
  * Kwon et al. (2021) Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi.  ASAM: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks.  In _Int. Conference on Machine Learning (ICML)_ , 2021.
  * Lawrence et al. (2022) Hannah Lawrence, Kristian Georgiev, Andrew Dienes, and Bobak T Kiani.  Implicit bias of linear equivariant networks.  In _Int. Conference on Machine Learning (ICML)_ , 2022.
  * Le and Jegelka (2022) Thien Le and Stefanie Jegelka.  Training invariances and the low-rank phenomenon: beyond linear networks.  In _Int. Conference on Learning Representations (ICLR)_ , 2022.
  * Li et al. (2022a) Zhiyuan Li, Tianhao Wang, and Sanjeev Arora.  What happens after sgd reaches zero loss?–a mathematical framework.  In _Int. Conference on Learning Representations (ICLR)_ , 2022a.
  * Li et al. (2022b) Zhiyuan Li, Tianhao Wang, and Dingli Yu.  Fast mixing of stochastic gradient descent with normalization and weight decay.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2022b.
  * Liang et al. (2019) Tengyuan Liang, Tomaso Poggio, Alexander Rakhlin, and James Stokes.  Fisher-rao metric, geometry, and complexity of neural networks.  In _Int. Conference on Artificial Intelligence and Statistics (AISTATS)_ , 2019.
  * Lim et al. (2023) Derek Lim, Joshua Robinson, Lingxiao Zhao, Tess Smidt, Suvrit Sra, Haggai Maron, and Stefanie Jegelka.  Sign and basis invariant networks for spectral graph representation learning.  In _Int. Conference on Learning Representations (ICLR)_ , 2023.
  * Liu et al. (2022a) Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You.  Towards efficient and scalable sharpness-aware minimization.  In _IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2022a.
  * Liu et al. (2022b) Yong Liu, Siqi Mai, Minhao Cheng, Xiangning Chen, Cho-Jui Hsieh, and Yang You.  Random sharpness-aware minimization.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2022b.
  * Long and Bartlett (2023) Philip M. Long and Peter L. Bartlett.  Sharpness-aware minimization and the edge of stability.  _arXiv preprint arXiv:2309.12488_ , 2023.
  * Lu et al. (2022) Peng Lu, Ivan Kobyzev, Mehdi Rezagholizadeh, Ahmad Rashid, Ali Ghodsi, and Philippe Langlais.  Improving generalization of pre-trained language models via stochastic weight averaging.  _arXiv preprint arXiv:2212.05956_ , 2022.
  * Lyu et al. (2022) Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora.  Understanding the generalization benefit of normalization layers: Sharpness reduction.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Madry et al. (2018) Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu.  Towards deep learning models resistant to adversarial attacks.  In _Int. Conference on Learning Representations (ICLR)_ , 2018.
  * Mi et al. (2022) Peng Mi, Li Shen, Tianhe Ren, Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao.  Make sharpness-aware minimization stronger: A sparsified perturbation approach.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Neyshabur et al. (2017) Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro.  Exploring generalization in deep learning.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2017.
  * Nitanda et al. (2023) Atsushi Nitanda, Ryuhei Kikuchi, and Shugo Maeda.  Parameter averaging for sgd stabilizes the implicit bias towards flat regions.  _arXiv preprint arXiv:2302.09376_ , 2023.
  * Qu et al. (2022) Zhe Qu, Xingyu Li, Rui Duan, Yao Liu, Bo Tang, and Zhuo Lu.  Generalized federated learning via sharpness aware minimization.  In _Int. Conference on Machine Learning (ICML)_ , 2022.
  * Shi et al. (2023) Yifan Shi, Yingqi Liu, Kang Wei, Li Shen, Xueqian Wang, and Dacheng Tao.  Make landscape flatter in differentially private federated learning.  In _IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2023.
  * Soudry et al. (2018) Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro.  The implicit bias of gradient descent on separable data.  _Journal of Machine Learning Research_ , 2018.
  * Sun et al. (2023) Hao Sun, Li Shen, Qihuang Zhong, Liang Ding, Shixiang Chen, Jingwei Sun, Jing Li, Guangzhong Sun, and Dacheng Tao.  AdaSAM: Boosting sharpness-aware minimization with adaptive learning rate and momentum for training deep neural networks.  _arXiv preprint arXiv:2303.00565_ , 2023.
  * Wang et al. (2023) Pengfei Wang, Zhaoxiang Zhang, Zhen Lei, and Lei Zhang.  Sharpness-aware gradient matching for domain generalization.  In _IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2023.
  * Wen et al. (2023a) Kaiyue Wen, Tengyu Ma, and Zhiyuan Li.  How does sharpness-aware minimization minimize sharpness?  In _Int. Conference on Learning Representations (ICLR)_ , 2023a.
  * Wen et al. (2023b) Kaiyue Wen, Tengyu Ma, and Zhiyuan Li.  Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization.  _arXiv preprint arXiv:2307.11007_ , 2023b.
  * Woodworth et al. (2020) Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro.  Kernel and rich regimes in overparametrized models.  In _Conference on Learning Theory (COLT)_ , 2020.
  * Xu et al. (2019) Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka.  How powerful are graph neural networks?  In _Int. Conference on Learning Representations (ICLR)_ , 2019.
  * Yao et al. (2020) Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney.  Pyhessian: Neural networks through the lens of the hessian.  In _2020 IEEE international conference on big data (Big data)_ , pages 581–590. IEEE, 2020.
  * Zaheer et al. (2017) Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola.  Deep sets.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2017.
  * Zhao et al. (2022) Yang Zhao, Hao Zhang, and Xiuyuan Hu.  Randomized sharpness-aware training for boosting computational efficiency in deep learning.  _arXiv preprint arXiv:2203.09962_ , 2022.
  * Zhong et al. (2022) Qihuang Zhong, Liang Ding, Li Shen, Peng Mi, Juhua Liu, Bo Du, and Dacheng Tao.  Improving sharpness-aware minimization with fisher mask for better generalization on language models.  _arXiv preprint arXiv:2210.05497_ , 2022.
  * Zhu et al. (2023) Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge.  Understanding edge-of-stability training dynamics with a minimalist example.  In _Int. Conference on Learning Representations (ICLR)_ , 2023.
  * Zhuang et al. (2022) Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvornek, Sekhar Tatikonda, James Duncan, and Ting Liu.  Surrogate gap minimization improves sharpness-aware training.  In _Int. Conference on Learning Representations (ICLR)_ , 2022.

##  Appendix A Additional Related Work

SAM can provide a strong regularization of the eigenvalues throughout the
learning trajectory (Agarwala and Dauphin, 2023). Bartlett et al. (2022) show
that the dynamics of SAM similar to GD on the spectral norm of Hessian.
Compagnoni et al. (2023) propose an SDE for modeling SAM, while Behdin and
Mazumder (2023) study the statistical benefits of SAM (see also (Li et al.,
2022a) for a general framework for the dynamics of SGD around the zero-loss
manifold). It is shown that SAM can reduce the feature rank (i.e., allowing
learning low-rank features) (Andriushchenko et al., 2023a). Blanc et al.
(2020) proved that SGD is implicitly biased toward minimizing the trace of
Hessian.

Kim et al. (2023) proposes a multi-step ascent approach to improve SAM, while
Mi et al. (2022) suggested sparsification of SAM. Zhuang et al. (2022) improve
SAM by changing the directions in the ascent step; their method is called
Surrogate Gap Guided Sharpness-Aware Minimization (GSAM) (see also (Behdin et
al., 2022)). Random smoothing-based SAM (R-SAM) is another SAM variant that is
proposed to reduce its computational complexity (Liu et al., 2022b) (see also
(Du et al., 2022; Liu et al., 2022a; Zhao et al., 2022; Sun et al., 2023) for
more). Adaptive SAM (ASAM) is proposed for applying SAM on scale-invariant
neural networks and has shown generalization benefits (Kwon et al., 2021). Li
et al. (2022b) also prove that scale-invariant loss functions allow faster
mixing in function spaces for neural networks. Lyu et al. (2022) show how
normalization can make GD reduce the sharpness via a continuous sharpness-
reduction flow. Liang et al. (2019) propose a capacity measure based on
information geometry for parameter invariances in overparameterized models
(for more on information geometry, see (Kim et al., 2022; Jang et al., 2022)).
Jiang et al. (2020) empirically compare different complexity measures for
overparameterized models. Keskar et al. (2017) show how a large batch yields
sharp minima but a small batch achieves flat minima.

Stochastic Weight Averaging (SWA) is another way to improve generalization and
it relies on finding wider minima by averaging multiple points along the
trajectory of SGD (Izmailov et al., 2018). (see e.g., (Lu et al., 2022) which
uses this method for language models). See (Kaddour et al., 2022) for the
empirical comparison between two popular flat-minima optimization approaches:
SWA and SAM.

Learning with group invariant architectures has recently gained a lot of
interest due to its applications in physics and biology; see e.g., deep sets
(Zaheer et al., 2017), Graph Neural Networks (GNNs) (Xu et al., 2019), and
also sign-flips for spectral data (Lim et al., 2023). These architectures are
all owing their practical success to their specific parameter invariance.

It is also worth mentioning that neural networks trained with large learning
rates often generalize better (the edge-of-stability regime); see (Arora et
al., 2022; Zhu et al., 2023; Long and Bartlett, 2023) for the theoretical
understanding of this phenomenon.

##  Appendix B Examples of $(\phi,\psi,\mu)$-Sharpness Measures

In this section, we prove various notions of sharpness can be achieved using
the proposed approach in this paper (Table 1). For the last row of Table 1, we
refer the reader to the proof of Theorem 1.

  * •

Trace. Let $\phi(t)=\psi(t)=t$, and note that

$\displaystyle S(x;\phi,\psi,\mu)=$
$\displaystyle\int\frac{1}{2}v^{t}\nabla^{2}L(x)vd\mu(v)$ (13)
$\displaystyle=\frac{1}{2}\mathbb{E}_{v\sim\mu}[v^{t}\nabla^{2}L(x)v],$ (14)

where $\mu$ is the uniform distribution over the $(d-1)$-sphere
$\mathbb{S}^{(d-1)}:=\\{x\in\mathbb{R}^{d}:\|x\|_{2}=1\\}$. Denote the entries
of $\nabla^{2}L(x)$ as $(\nabla^{2}L(x))_{i,j}$. Then, by the linearity of
expectation

$\displaystyle\mathbb{E}_{v\sim\mu}[v^{t}\nabla^{2}L(x)]$
$\displaystyle=\sum_{i,j=1}^{d}(\nabla^{2}L(x))_{i,j}\mathbb{E}[v_{i}v_{j}]=\sum_{i=1}^{d}\frac{1}{d}(\nabla^{2}L(x))_{i,i}=\frac{1}{d}\operatorname*{tr}(\nabla^{2}L(x)),$
(15)

since $\mathbb{E}[v_{i}v_{j}]=\frac{1}{d}\delta_{i,j}$, where $\delta_{i,j}$
denotes the Knocker delta function.

  * •

Determinant. To achieve the determinant, we choose $\phi(t)=(2\pi)^{d}/t^{2}$
and $\psi(t)=\exp(-t)$. Then,

$\displaystyle S(x;\phi,\psi,\mu)=$
$\displaystyle(2\pi)^{d}\Big{(}\int\exp(-\frac{1}{2}v^{t}\nabla^{2}L(x)v)dv\Big{)}^{-2},$
(16)

where $dv$ denotes the Lebesgue measure. However, using the multivariate
Gaussian integral, we have

$\displaystyle\int\exp\Big{(}-\frac{1}{2}v^{t}\nabla^{2}L(x)v\Big{)}dv=(2\pi)^{d/2}\det(\nabla^{2}L(x))^{-1/2}.$
(17)

Replacing this intro the definition of $S(x;\phi,\psi,\mu)$ gives the desired
result.

  * •

Polynomials of eigenvalues. First assume that $\psi(t)=t^{n}$ for some $n\geq
0$. Then, for any function $\phi(t)$,

$\displaystyle S(x;\phi,\psi,\mu)$
$\displaystyle=\phi\Big{(}\int\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}^{n}d\mu(v)\Big{)}$
(18)
$\displaystyle=\phi\Big{(}\mathbb{E}_{v\sim\mu}\Big{[}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}^{n}\Big{]}\Big{)},$
(19)

where $\mu$ is the uniform distribution over the $(d-1)$-sphere
$\mathbb{S}^{(d-1)}$. Since $\nabla^{2}L(x)$ is a symmetric matrix, we can
find an orthogonal matrix $Q$ such that $\nabla^{2}L(x)=Q^{t}DQ$, where $D$ is
a diagonal matrix with diagonal entries
$\lambda_{1},\lambda_{2},\ldots,\lambda_{d}$. Now we write
$\big{(}v^{t}\nabla^{2}L(x)v\big{)}^{n}=\big{(}v^{t}Q^{t}DQv\big{)}^{n}$. But
$Qv$ is distributed uniformly over the $(d-1)$-sphere $\mathbb{S}^{(d-1)}$,
similar to $v$. Thus, we conclude

$\displaystyle S(x;\phi,\psi,\mu)$
$\displaystyle=\phi\Big{(}\mathbb{E}_{v\sim\mu}\Big{[}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}^{n}\Big{]}\Big{)}$
(20)
$\displaystyle=\phi\Big{(}\mathbb{E}_{v\sim\mu}\Big{[}\big{(}\frac{1}{2}\sum_{i=1}^{d}\lambda_{i}v_{i}^{2}\big{)}^{n}\Big{]}\Big{)}.$
(21)

Define

$\displaystyle
q(\lambda_{1},\lambda_{2},\ldots,\lambda_{d}):=\mathbb{E}_{v\sim\mu}\Big{[}\big{(}\sum_{i=1}^{d}\frac{1}{2}\lambda_{i}v_{i}^{2}\big{)}^{n}\Big{]},$
(22)

which is clearly a polynomial function (by the linearity of expectation).

Note that the above computation is still valid if we replace the uniform
distribution on hypersphere with the Gaussian multivariate distribution with
identity covariance $\mathcal{N}(0,I_{d})$. Indeed, let us compute this
polynomial for $n=2$ with Gaussian distribution. Note that

$\displaystyle q(\lambda_{1},\lambda_{2},\ldots,\lambda_{d})$
$\displaystyle=\mathbb{E}_{v\sim\mu}\Big{[}\big{(}\sum_{i=1}^{d}\frac{1}{2}\lambda_{i}v_{i}^{2}\big{)}^{2}\Big{]}=\frac{1}{4}\sum_{i=1}^{d}\lambda_{i}^{2}\mathbb{E}[Z^{4}]+\sum_{i\neq
j}\frac{1}{4}\lambda_{i}\lambda_{j}(\mathbb{E}[Z^{2}])^{2}$ (23)
$\displaystyle=\frac{3}{4}\sum_{i=1}^{d}\lambda_{i}^{2}+\frac{1}{4}\sum_{i\neq
j}\lambda_{i}\lambda_{j},$ (24)

where $Z$ is a zero-mean Gaussian random variable with unit variance, and note
that $E[Z^{2}]=1$ and $E[Z^{4}]=3$.

Now if we take $m=2$, and $\psi(t)=(t,t^{2})$, with
$\mu=\mathcal{N}(0,I_{d})\otimes\mathcal{N}(0,I_{d})$, we have that

$\displaystyle\int\psi\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)=\Big{(}\frac{1}{2}\sum_{i=1}^{d}\lambda_{i},\frac{3}{4}\sum_{i=1}^{d}\lambda_{i}^{2}+\frac{1}{4}\sum_{i\neq
j}\lambda_{i}\lambda_{j}\Big{)}.$ (25)

Finally, by taking $\phi(t_{1},t_{2})=2(t_{2}-t_{1}^{2})$, we obtain

$\displaystyle\phi\Big{(}\int\psi\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)\Big{)}=\sum_{i=1}^{d}\lambda_{i}^{2}.$
(26)

##  Appendix C Proof of Theorem 1

See 1

###### Proof.

We explicitly construct the (Borel) probability measure $\mu$ and the function
$\psi:\mathbb{R}^{m}\to\mathbb{R}$. Let us take $m=d$ to prove the
universality theorem, while we believe lower $m$ should be enough for specific
practical sharpness measures.

Indeed, let us consider $\mu$ to be the multivariate Gaussian probability
measure with identity covariance matrix. Also, define a (parameterized)
function $\psi_{\sigma}(t)=\exp(\sigma t)$, for some $\sigma$ to be set later.
We are interested to compute the following quantity:

$\displaystyle\int\psi_{\sigma}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v).$
(27)

Note that $\mu$ is the standard Gaussian probability measure on
$\mathbb{R}^{d}$, and specifically, it’s invariant under the action of the
orthogonal matrices. Indeed, there exists an orthogonal matrix $Q$ such that
$\nabla^{2}L(x)=Q^{t}\Lambda Q$, where $\Lambda$ is a diagonal matrix with
diagonal entries $\lambda_{1},\lambda_{2},\ldots,\lambda_{d}$. Observe that
$u:=Qv$ is also distributed according to the Gaussian probability distribution
with the identity covariance matrix on $\mathbb{R}^{d}$, similar to $v$. Now
we write

$\displaystyle\int\psi_{\sigma}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)$
$\displaystyle=\int\psi_{\sigma}\big{(}\frac{1}{2}v^{t}Q^{t}\Lambda
Qv\big{)}d\mu(v)$ (28)
$\displaystyle=\int\psi_{\sigma}\big{(}\frac{1}{2}u^{t}\Lambda
u\big{)}d\mu(u)$ (29)
$\displaystyle=\int\psi_{\sigma}\Big{(}\frac{1}{2}\sum_{i=1}^{d}\lambda_{i}u_{i}^{2}\Big{)}d\mu(u)$
(30)
$\displaystyle=\int(2\pi)^{-d/2}\psi_{\sigma}\Big{(}\frac{1}{2}\sum_{i=1}^{d}\lambda_{i}u_{i}^{2}\Big{)}\exp\Big{(}-\frac{1}{2}\sum_{i=1}^{d}u_{i}^{2}\Big{)}du$
(31)
$\displaystyle=\int(2\pi)^{-d/2}\exp\Big{(}\frac{1}{2}\sum_{i=1}^{d}\sigma\lambda_{i}u_{i}^{2}\Big{)}\exp\Big{(}-\frac{1}{2}\sum_{i=1}^{d}u_{i}^{2}\Big{)}du$
(32)
$\displaystyle=\int(2\pi)^{-d/2}\exp\Big{(}\sum_{i=1}^{d}\frac{1}{2}(\sigma\lambda_{i}-1)u_{i}^{2}\Big{)}du,$
(33)

where $du$ denotes the Lebesgue measure on $\mathbb{R}^{d}$. Now to compute
the integral, note that $du=du_{1}\times du_{2}\times\ldots\times du_{d}$ is a
product measure and the integrand also takes on a product form; thus,

$\displaystyle\int\psi_{\sigma}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)$
$\displaystyle=\int(2\pi)^{-d/2}\exp\Big{(}\sum_{i=1}^{d}\frac{1}{2}(\sigma\lambda_{i}-1)u_{i}^{2}\Big{)}du$
(34)
$\displaystyle=\prod_{i=1}^{d}\int(2\pi)^{-1/2}\exp\Big{(}\frac{1}{2}(\sigma\lambda_{i}-1)u_{i}^{2}\Big{)}du_{i}$
(35)
$\displaystyle\overset{(a)}{=}\prod_{i=1}^{d}\frac{1}{\sqrt{1-\sigma\lambda_{i}}}\underbrace{\int\sqrt{\frac{1-\sigma\lambda_{i}}{2\pi}}\exp\Big{(}\frac{1}{2}(\sigma\lambda_{i}-1)u_{i}^{2}\Big{)}du_{i}}_{=1}$
(36) $\displaystyle=\prod_{i=1}^{d}\frac{1}{\sqrt{1-\sigma\lambda_{i}}},$ (37)

where (a) holds by the Gaussian integral identities.

Note that to calculate the integral above, we assumed that
$1-\sigma\lambda_{i}>0$ for any $i\in[d]$. This is equivalent to having

$\displaystyle\max_{\lambda_{i}<0}\lambda^{-1}_{i}<\sigma<\min_{\lambda_{i}>0}\lambda^{-1}_{i}.$
(38)

Now, due to the assumption in the theorem, we study the target sharpness
measure $S(\lambda_{1},\lambda_{2},\ldots,\lambda_{d})$ only on a compact
domain $\mathcal{A}\subseteq\mathbb{R}^{d}$, and this means that there exists
an open interval $I=(-\epsilon,\epsilon)$ with

$\displaystyle\epsilon:=\min_{\big{(}\lambda_{1},\lambda_{2},\ldots,\lambda_{d}\big{)}^{t}\in\mathcal{A}}\leavevmode\nobreak\
\min_{i}|\lambda_{i}|^{-1},$ (39)

such that the above integral is well-defined and finite for all $\sigma\in I$.

Let us define the function $\tilde{\phi}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ as
follows:

$\displaystyle\tilde{\phi}(t_{1},t_{2},\ldots,t_{d}):=(t_{1}^{-2},t_{2}^{-2},\ldots,t_{d}^{-2}).$
(40)

Finally, consider the following polynomial in one variable with degree $d$:

$\displaystyle
p(x):=(1-\lambda_{1}x)\times(1-\lambda_{2}x)\times\cdots\times(1-\lambda_{d}x).$
(41)

######  Claim 1.

For any $\sigma_{i}\in I$, $i\in[d]$, we have

$\displaystyle\tilde{\phi}\Big{(}\int\psi_{\sigma_{1}}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v),\int\psi_{\sigma_{2}}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v),\ldots,$
$\displaystyle\int\psi_{\sigma_{d}}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)\Big{)}$
(42)
$\displaystyle=\Big{(}p(\sigma_{1}),p(\sigma_{2}),\ldots,p(\sigma_{d})\Big{)}.$
(43)

The above claim simply follows from the integral we calculated before.

Now we are ready to complete the proof. Choose arbitrary non-zero distinct
$\sigma_{i}\in I,i\in[d]$, and note that having access to
$\big{(}p(\sigma_{1}),p(\sigma_{2}),\ldots,p(\sigma_{d})\big{)}$ is enough to
recover all the eigenvalues. Indeed, assume that
$p(x)=p_{0}+p_{1}x+p_{2}x^{2}+\ldots+p_{d}x^{d}$ and note that $p(0)=1$ by
definition. Let also
$V(0,\sigma_{1},\sigma_{2},\ldots,\sigma_{d})\in\mathbb{R}^{d\times d}$ denote
a Vandermonde matrix of order $d+1$, which is provably invertible by
definition, and note that

$\displaystyle\big{(}p(\sigma_{1}),p(\sigma_{2}),$
$\displaystyle\ldots,p(\sigma_{d})\big{)}^{t}=V(0,\sigma_{1},\sigma_{2},\ldots,\sigma_{d})\times\big{(}p_{0},p_{1},\ldots,p_{d+1}\big{)}^{t}$
(44)
$\displaystyle\implies\big{(}p_{0},p_{1},\ldots,p_{d+1}\big{)}^{t}=V(0,\sigma_{1},\sigma_{2},\ldots,\sigma_{d})^{-1}\times\big{(}p(\sigma_{1}),p(\sigma_{2}),\ldots,p(\sigma_{d})\big{)}^{t}.$
(45)

Indeed, this shows that having access to the vector
$\big{(}p(\sigma_{1}),p(\sigma_{2}),\ldots,p(\sigma_{d})\big{)}^{t}$ is enough
to reconstruct the polynomial $p(x)=p_{0}+p_{1}x+\ldots+p_{d}x^{d}$. Having
access to this polynomial is equivalent to having access to its roots, so one
can find a continuous function $\phi_{1}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ such
that

$\displaystyle\big{(}$
$\displaystyle\lambda_{1},\lambda_{2},\ldots,\lambda_{d}\big{)}^{t}$
$\displaystyle=\phi_{1}\circ\tilde{\phi}\Big{(}\int\psi_{\sigma_{1}}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v),\int\psi_{\sigma_{2}}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v),\ldots,\int\psi_{\sigma_{d}}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)\Big{)}.$

Since the sharpness measure $S(\lambda_{1},\lambda_{2},\ldots,\lambda_{d})$ is
a continuous function of its coordinates, we conclude that

$\displaystyle S(\lambda_{1},\lambda_{2},\ldots,\lambda_{d}\big{)}$
$\displaystyle=S\circ\phi_{1}\circ\tilde{\phi}\Big{(}\int\psi_{\sigma_{1}}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v),\int\psi_{\sigma_{2}}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v),\ldots,\int\psi_{\sigma_{d}}\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)\Big{)}.$

Now to complete the proof, we define a continuous function
$\psi:\mathbb{R}\to\mathbb{R}^{d}$ as
$\psi=\big{(}\psi_{\sigma_{1}},\psi_{\sigma_{2}},\ldots,\psi_{\sigma_{d}}\big{)}^{t}$,
and a continuous function $\phi:\mathbb{R}^{d}\to\mathbb{R}$ as
$\phi=S\circ\phi_{1}\circ\tilde{\phi}$, and observe that
$S(\lambda_{1},\lambda_{2}\ldots,\lambda_{d})=S(x;\phi,\psi,\mu)$ for any
$x\in\mathcal{A}$. This completes the proof. ∎

##  Appendix D Proof of Theorem 2

See 2

###### Proof.

We explicitly construct a set of functions/probability measures to achieve the
desired representation. Indeed, let’s take $m=d(d+1)/2$ and consider the
following Dirac measures: $\mu_{i}=\delta_{e_{i}},i\in[d]$, and also
$\mu_{ij}=\delta_{e_{i}+e_{j}}$, for any $i,j\in[d]$ such that $i<j$. Here,
$e_{i}$ denotes the unit vector in the $i$th coordinate in $\mathbb{R}^{d}$.
Now, note that we have

$\displaystyle(\nabla^{2}L(x))_{i,i}=\int v^{t}\nabla^{2}L(x)vd\mu_{i}(v)$
(46)

for any $i\in[d]$, and

$\displaystyle
2(\nabla^{2}L(x))_{i,j}+(\nabla^{2}L(x))_{i,i}+(\nabla^{2}L(x))_{j,j}=\int
v^{t}\nabla^{2}L(x)vd\mu_{i,j}(v),$ (47)

for any $i,j\in[d]$ such that $i<j$. The above system of linear equations has
clearly a unique solution, as the Hessian matrix is symmetric. This means
that, similar to the proof of Theorem 1, one can find continuous functions
$\psi:\mathbb{R}\to\mathbb{R}^{m}$ and $\phi:\mathbb{R}^{m}\to R$, along with
$m$ constructed probability measures such that
$S(\nabla^{2}L(x))=S(x;\phi,\psi,\mu)$ for any $x\in\mathbb{R}^{d}$, where
$\mu=\mu_{1}\otimes\mu_{2}\otimes\ldots\otimes\mu_{m}$ is a product
probability measure. ∎

##  Appendix E Proof of Theorem 3

See 3

###### Proof.

For simplicity, we assume that $m=1$. The general proof for $m>1$ follows with
a similar argument to this special case. Define an open set
$U\subseteq\mathbb{R}^{d}$ as follows:

$\displaystyle U:=\big{\\{}x\in\mathbb{R}^{d}:\|\nabla
L(x)\|_{2}<\rho^{2}\big{\\}}.$ (48)

Note that this set contain the zero-loss manifold, i.e., $\Gamma\subseteq U$.
We study the behavior of the loss function on this open set.

Let us denote the sharpness term in the loss function $L_{(\phi,\psi,\mu)}(x)$
by $R_{\rho}(x)$:

$\displaystyle L_{(\phi,\psi,\mu)}(x)$
$\displaystyle=L(x)+\rho^{2}R_{\rho}(x)$ (49)
$\displaystyle=L(x)+\rho^{2}\phi\Big{(}\int\psi\Big{(}\frac{1}{\rho^{2}}\big{(}L(x+\rho
v)-L(x)\big{)}\Big{)}d\mu(v)\Big{)}.$ (50)

We first study the convergence of $R_{\rho}(x)$ to the corresponding sharpness
measure $S(x;\phi,\psi,\mu)$. Fix any point $x\in U$ and note that using
Taylor’s theorem for the function $L:\mathbb{R}^{d}\to\mathbb{R}$ and for any
$v\in\mathbb{R}^{d}$, one has

$\displaystyle\frac{1}{\rho^{2}}\big{(}L(x+\rho v)-L(x)\big{)}$
$\displaystyle=\frac{1}{\rho^{2}}\Big{(}\rho\langle\nabla
L(x),v\rangle+\frac{1}{2}\rho^{2}v^{t}\nabla^{2}L(x)v+O_{x}(\rho^{3}\|v\|_{2}^{3}\times\|v\|_{2}^{-1})\Big{)}$
(51) $\displaystyle=\rho^{-1}\langle\nabla
L(x),v\rangle+\frac{1}{2}v^{t}\nabla^{2}L(x)v+O_{x}(\rho\|v\|_{2}^{2}\times\|v\|_{2}^{-1}),$
(52)

where in above we used the fact that $L(x)$ is third-order continuously
differentiable and its third-order derivative satisfies the estimate

$\displaystyle\max_{i,j,k\in\\{1,2,3\\}}|\partial_{i}\partial_{j}\partial_{k}L(v)|=O(\|v\|^{-1}),$
(53)

for $v\in\mathbb{R}^{d}$ as $\|v\|_{2}\to\infty$.

Note that using the assumption $x\in U$, we have that

$\displaystyle|\rho^{-1}\langle\nabla L(x),v\rangle|\leq\rho^{-1}\|\nabla
L(x)\|_{2}\|v\|_{2}<\rho\|v\|_{2}.$ (54)

Thus, we have

$\displaystyle\frac{1}{\rho^{2}}\big{(}L(x+\rho
v)-L(x)\big{)}=\frac{1}{2}v^{t}\nabla^{2}L(x)v+O_{x}(\rho(\|v\|_{2}^{2}+\|v\|_{2})).$
(55)

Note that we study the above approximation only for $x\in U$, and for small
enough $\rho$, we know that $U$ is a precompact set. Therefore, we drop the
dependence on $x$ in the error term above.

Now using the above approximation, we have

$\displaystyle\int\psi\Big{(}\frac{1}{\rho^{2}}\big{(}L(x+\rho
v)-L(x)\big{)}\Big{)}d\mu(v)=\int\psi\Big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v+O(\rho(\|v\|_{2}^{2}+\|v\|_{2}))\Big{)}d\mu(v).$
(56)

Let us use Taylor’s theorem for the function $\psi$ and write

$\displaystyle\psi\Big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v+O(\rho(\|v\|_{2}^{2}+\|v\|_{2}))\Big{)}=\psi\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}+\rho\times
O(\tilde{\psi}(v)(\|v\|_{2}^{2}+\|v\|_{2})),$ (57)

where

$\displaystyle\tilde{\psi}(v):=\max_{|t|\leq\|v\|_{2}}|\psi^{\prime}(Ct^{2})|,$
(58)

and $C$ is a constant, and it’s big enough to absorb the quadratic growth of
$\nabla^{2}L(x)$; i.e.,

$\displaystyle
C>\max_{x\in\mathcal{X}}\max_{i\in[d]}|\lambda_{i}(\nabla^{2}L(x))|.$ (59)

Therefore, we conclude that

$\displaystyle\int\psi\Big{(}\frac{1}{\rho^{2}}\big{(}L(x+\rho
v)-L(x)\big{)}\Big{)}d\mu(v)$
$\displaystyle=\int\psi\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)+\rho\times
O\Big{(}\int\tilde{\psi}(v)(\|v\|_{2}^{2}+\|v\|_{2})d\mu(v)\Big{)}$
$\displaystyle=\int\psi\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)+O(\rho),$

by the assumption. This allows us to conclude that

$\displaystyle R_{\rho}(x)$
$\displaystyle=\phi\Big{(}\int\psi\Big{(}\frac{1}{\rho^{2}}\big{(}L(x+\rho
v)-L(x)\big{)}\Big{)}d\mu(v)\Big{)}$ (60)
$\displaystyle=\phi\Big{(}\int\psi\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)\Big{)}+O(\rho)$
(61) $\displaystyle=S(x;\phi,\psi,\mu)+O(\rho),$ (62)

again, by assuming that

$\displaystyle\max_{x\in\mathcal{X}}\phi^{\prime}\Big{(}\int\psi\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)v\big{)}d\mu(v)\big{)}<\infty,$
(63)

which holds by the compactness of $\mathcal{X}$, and also using
$U\subseteq\mathcal{X}$.

Now according to the assumption, for some $u\in U$, we have

$\displaystyle L(u)+\rho^{2}R_{\rho}(u)-\inf_{x\in
U}\Big{(}L(x)+\rho^{2}R_{\rho}(x)\Big{)}\leq\Delta\rho^{2},$ (64)

for some optimally gap $\Delta$. Using the following proven approximation

$\displaystyle R_{\rho}(x)=S(x;\phi,\psi,\mu)+O(\rho),$ (65)

we conclude that

$\displaystyle L(u)+\rho^{2}S(u;\phi,\psi,\mu)-\inf_{x\in
U}\Big{(}L(x)+\rho^{2}S(x;\phi,\psi,\mu)\Big{)}\leq(\Delta+O(\rho))\rho^{2}.$
(66)

Now by proof by contradiction, assume that

$\displaystyle L(u)\geq\inf_{x\in U}L(x)+(\Delta+\delta)\rho^{2},$ (67)

for some strictly positive $\delta$, as $\rho\to 0^{+}$. Note that $\inf_{x\in
U}L(x)=0$ as $\Gamma\subseteq U$. Thus, we can conclude that

$\displaystyle\rho^{2}S(u;\phi,\psi,\mu)+(\Delta+\delta)\rho^{2}$
$\displaystyle\leq L(u)+\rho^{2}S(u;\phi,\psi,\mu)$ (68)
$\displaystyle\leq\inf_{x\in
U}\Big{(}L(x)+\rho^{2}S(x;\phi,\psi,\mu)\Big{)}+(\Delta+O(\rho))\rho^{2}$ (69)
$\displaystyle\leq\rho^{2}\inf_{x\in\Gamma}S(x;\phi,\psi,\mu)+(\Delta+O(\rho))\rho^{2},$
(70)

since $\Gamma\subseteq U$. This shows that

$\displaystyle
S(u;\phi,\psi,\mu)\leq\inf_{x\in\Gamma}S(x;\phi,\psi,\mu)-\delta+O(\rho).$
(71)

This must hold for as $\rho\to 0^{+}$. However, as $\rho\to 0^{+}$, we have
that $U_{\rho}\to\Gamma$. This means that

$\displaystyle
S(u;\phi,\psi,\mu)\leq\inf_{x\in\Gamma}S(x;\phi,\psi,\mu)-\delta,$ (72)

for some $u\in\Gamma$, which is a contradiction. This shows that

$\displaystyle L(u)\leq\inf_{x\in U}L(x)+(\Delta+o(1))\rho^{2}.$ (73)

Also, to prove the next part of the theorem, for any $u\in U$ satisfying the
assumptions, similarly we can show

$\displaystyle\rho^{2}S(u;\phi,\psi,\mu)$ $\displaystyle\leq
L(u)+\rho^{2}S(u;\phi,\psi,\mu)$ (74) $\displaystyle=\inf_{x\in
U}\Big{(}L(x)+\rho^{2}S(x;\phi,\psi,\mu)\Big{)}+(\Delta+O(\rho))\rho^{2}$ (75)
$\displaystyle\leq\rho^{2}\inf_{x\in\Gamma}S(x;\phi,\psi,\mu)+(\Delta+O(\rho))\rho^{2},$
(76)

which implies that

$\displaystyle
S(u;\phi,\psi,\mu)\leq\inf_{x\in\Gamma}S(x;\phi,\psi,\mu)+(\Delta+O(\rho)).$
(77)

The proof is thus complete.

∎

##  Appendix F Proof of Theorem 4

See 4

###### Proof.

Let $D\in\mathbb{R}^{d\times d}$ be an arbitrary diagonal matrix. Then,

$\displaystyle
S(Dx;\phi,\psi,\mu)=\phi\Big{(}\int\psi\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)\Big{|}_{Dx}v\big{)}d\mu(v)\Big{)}.$
(78)

But note that by the assumption

$\displaystyle
L(x)=L(Dx)\implies\nabla^{2}L(x)=D^{t}\nabla^{2}L(x)\Big{|}_{Dx}D.$ (79)

Therefore,

$\displaystyle
S(Dx;\phi,\psi,\mu)=\phi\Big{(}\int\psi\big{(}\frac{1}{2}v^{t}D^{-1}\nabla^{2}L(x)D^{-1}v\big{)}d\mu(v)\Big{)}.$
(80)

Now define a new variable $u:=D^{-1}v$. Then,

$\displaystyle d\mu(v)$
$\displaystyle=f(\prod_{i=1}^{d}v_{i})\prod_{i=1}^{d}dv_{i}=f(\prod_{i=1}^{d}D_{i,i}\prod_{i=1}^{d}u_{i})\prod_{i=1}^{d}D_{i,i}\prod_{i=1}^{d}du_{i}$
(81)
$\displaystyle=f(\det(D)\prod_{i=1}^{d}u_{i})\det(D)\prod_{i=1}^{d}du_{i}$
(82) $\displaystyle=f(\prod_{i=1}^{d}u_{i})\prod_{i=1}^{d}du_{i}.$ (83)

Therefore, we conclude that

$\displaystyle
S(Dx;\phi,\psi,\mu)=\phi\Big{(}\int\psi\big{(}\frac{1}{2}u^{t}\nabla^{2}L(x)u\big{)}d\mu(u)\Big{)}=S(x;\phi,\psi,\mu),$
(84)

and this completes the proof.

∎

######  Lemma 1.

For any scale-invariant measure $\mu$ that is absolutely continuous with
respect to the Lebesgue measure on $\mathbb{R}^{d}$, one has

$d\mu(x)=f\big{(}\prod_{i=1}^{d}x_{i}\big{)}\prod_{i=1}^{d}dx_{i}.$

###### Proof.

Any measure $\mu$ that is absolutely continuous with respect to the Lebesgue
measure on $\mathbb{R}^{d}$ can be written as

$d\mu(x_{1},x_{2},\dots,x_{d})=\tilde{f}(x_{1},x_{2},\dots,x_{d})\prod_{i=1}^{d}dx_{i}.$

For any non-zero choice of $x_{1},x_{2},\dots,x_{d}$, by scale invariance
property of $d\mu$,

$\displaystyle
d\mu(x_{1},x_{2},\dots,x_{d})=d\mu(1,1,\dots,1,\prod_{i=1}^{d}x_{i})$
$\displaystyle\Rightarrow$
$\displaystyle\tilde{f}(x_{1},x_{2},\dots,x_{d})\prod_{i=1}^{d}dx_{i}=\tilde{f}(1,1,\dots,1,\prod_{i=1}^{d}x_{i})\prod_{i=1}^{d}dx_{i}$
$\displaystyle\Rightarrow$
$\displaystyle\tilde{f}(x_{1},x_{2},\dots,x_{d})=\tilde{f}(1,1,\dots,1,\prod_{i=1}^{d}x_{i}).$

Now, it suffices to choose
$f(\prod_{i=1}^{d}x_{i}):=\tilde{f}(1,1,\dots,1,\prod_{i=1}^{d}x_{i})$ almost
surely and the proof is concluded. ∎

##  Appendix G Proof of Theorem 5

See 5

###### Proof.

We start by evaluating $S(A_{g}x;\phi,\psi,\mu)$.

$\displaystyle
S(Dx;\phi,\psi,\mu)=\phi\Big{(}\int\psi\big{(}\frac{1}{2}v^{t}\nabla^{2}L(x)\Big{|}_{A_{g}x}v\big{)}d\mu(v)\Big{)}.$
(85)

But again here, note that by the assumption

$\displaystyle
L(x)=L(A_{g}x)\implies\nabla^{2}L(x)=A_{g}^{t}\nabla^{2}L(x)\Big{|}_{A_{g}x}A_{g}.$
(86)

Therefore,

$\displaystyle
S(A_{g}x;\phi,\psi,\mu)=\phi\Big{(}\int\psi\big{(}\frac{1}{2}v^{t}A_{g}^{-1}\nabla^{2}L(x)A_{g}^{-1}v\big{)}d\mu(v)\Big{)}.$
(87)

Now define a new variable $u:=A_{g}^{-1}v$. Therefore, we conclude that

$\displaystyle
S(A_{g}x;\phi,\psi,\mu)=\phi\Big{(}\int\psi\big{(}\frac{1}{2}u^{t}\nabla^{2}L(x)u\big{)}d\mu(u)\Big{)}=S(x;\phi,\psi,\mu),$
(88)

and this completes the proof.

∎

##  Appendix H $(\phi,\psi,\mu)$-Sharpness-Aware Minimization Algorithm

To propose an algorithm for the general case (i.e., arbitrary $m$), we compute
the gradient of

$\displaystyle\tilde{R}_{\rho}(x)\coloneqq\phi\Big{(}$
$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\psi_{1}\Big{(}\frac{1}{\rho^{2}}\big{(}L(x+\rho
v_{i,1})-L(x)\big{)}\Big{)},$ (89)
$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\psi_{2}\Big{(}\frac{1}{\rho^{2}}\big{(}L(x+\rho
v_{i,2})-L(x)\big{)}\Big{)},$ (90) $\displaystyle\hskip 56.9055pt\ldots$ (91)
$\displaystyle\frac{1}{n}\sum_{i=1}^{n}\psi_{m}\Big{(}\frac{1}{\rho^{2}}\big{(}L(x+\rho
v_{i,m})-L(x)\big{)}\Big{)}\Big{)},$ (92)

where $v_{i,\ell}\overset{\text{i.i.d.}}{\sim}\mu_{\ell}$ for each
$\ell\in[m]$. Note that $\psi=(\psi_{1},\psi_{2},\ldots,\psi_{m})^{t}$ for
some scalar functions $\psi_{\ell}$, $\ell\in[m]$. Let $\partial_{\ell}\phi$
denote partial derivatives of the function $\phi:\mathbb{R}^{m}\to\mathbb{R}$,
for any $\ell\in[m]$. Then,

$\displaystyle\rho^{2}\nabla\tilde{R}_{\rho}(x)$
$\displaystyle=\sum_{\ell=1}^{m}\partial_{\ell}\phi\Big{(}\sum_{i=1}^{n}\frac{1}{n}\psi_{\ell}\Big{(}\frac{1}{\rho^{2}}\big{(}L(x_{t}+\rho
v_{i,\ell})-L(x_{t})\big{)}\Big{)}\Big{)}$ (93) $\displaystyle\hskip
28.45274pt\times\sum_{i=1}^{n}\frac{1}{n}\Big{\\{}\psi_{\ell}^{\prime}\Big{(}\frac{1}{\rho^{2}}\big{(}L(x_{t}+\rho
v_{i,\ell})-L(x_{t})\big{)}\Big{)}\times\Big{(}\nabla L(x_{t}+\rho
v_{i,\ell})-\nabla L(x_{t})\Big{)}\Big{\\}},$ (94)

and this leads to Algorithm 3.

Algorithm 3 $(\phi,\psi,\mu)$-Sharpness-Aware Minimization Algorithm (with
arbitrary $m$)

Input: The triplet $(\phi,\psi,\mu)$, training loss $L(x)$, step size $\eta$,
perturbation parameter $\rho$, number of samples $n$,

Output: Model parameters $x_{t}$ trained with $(\phi,\psi,\mu)$-sharpness-
aware minimization algorithm

Initialization: $x\leftarrow x_{0}$ and $t\leftarrow 0$

while 1 do

Sample $v_{i,\ell}\overset{\text{i.i.d.}}{\sim}\mu_{\ell}$, for any $i\in[n]$
and $\ell\in[m]$

Compute the following:  $\displaystyle g_{t}$ $\displaystyle=\nabla
L(x_{t})+\sum_{\ell=1}^{m}\partial_{\ell}\phi\Big{(}\sum_{i=1}^{n}\frac{1}{n}\psi_{\ell}\Big{(}\frac{1}{\rho^{2}}\big{(}L(x_{t}+\rho
v_{i,\ell})-L(x_{t})\big{)}\Big{)}\Big{)}$ $\displaystyle\hskip
56.9055pt\times\sum_{i=1}^{n}\frac{1}{n}\Big{\\{}\psi_{\ell}^{\prime}\Big{(}\frac{1}{\rho^{2}}\big{(}L(x_{t}+\rho
v_{i,\ell})-L(x_{t})\big{)}\Big{)}\times\Big{(}\nabla L(x_{t}+\rho
v_{i,\ell})-\nabla L(x_{t})\Big{)}\Big{\\}}.$

Update the parameters: $x_{t+1}=x_{t}-\eta g_{t}$

$t\leftarrow t+1$

end while

##  Appendix I Experiments

###  I.1 Experimental Details

We now describe experimental details that were omitted from the main text.

For CIFAR10 and CIFAR100, we apply random crops and random horizontal flips.
We use a momentum term of 0.9 for all datasets and a weight decay of 5e-4 for
CIFAR10 and SVHN and 1e-3 for CIFAR100. We use batch size 128 and train for
200 epochs. We use a multi-step schedule where the learning rate is initially
0.1 and decays by a multiplicative factor of 0.1 every 50 epochs. We run each
experiment with four different random seeds to assess statistical
significance. We use 1280 training examples and 100 noise samples to estimate
the Frobenius norm and trace via Hessian-vector products. We set $\rho$ to 1.0
for Det-SAM, 0.01 for Trace-SAM, and sweep it in {0.005, 0.01} for Frob-SAM.
For Det-SAM and Trace-SAM we sweep $\lambda$ in {0.01, 0.1, 1.0} and set
$n=1$. For Frob-SAM, we sweep $\lambda$ in {0.0001, 0.001, 0.005, 0.01, 0.05,
0.1} and set $n=2$. The hyper-parameters selected for each setting is given in
Table 4 and Table 5.

We use the PyHessian library (Yao et al., 2020) to estimate the trace of the
Hessian. Adapting this library, we estimate the Frobenius norm (squared) as
$\frac{1}{k}\sum_{i=1}^{k}\|Hz_{i}\|_{2}^{2}$, where $H$ is the Hessian matrix
and $z_{i}\sim\mathcal{N}(0,I_{d})$.

| CIFAR10 | CIFAR100 | SVHN | CIFAR10-S | CIFAR100-S | SVHN-S | CIFAR10-C | CIFAR100-C | SVHN-C
---|---|---|---|---|---|---|---|---|---
Frob-SAM | 5e-3 | 1e-4 | 5e-3 | 1e-4 | 1e-4 | 1e-3 | 0.01 | 0.01 | 1e-3
Trace-SAM | 1 | 0.1 | 0.1 | 0.01 | 0.01 | 1 | 1 | 0.1 | 0.01
Det-SAM | 1 | 0.1 | 0.1 | 0.01 | 1 | 1 | 0.01 | 0.1 | 0.1
Table 4: Choice of regularizer weight $\lambda$ for different datasets. | CIFAR10 | CIFAR100 | SVHN | CIFAR10-S | CIFAR100-S | SVHN-S | CIFAR10-C | CIFAR100-C | SVHN-C
---|---|---|---|---|---|---|---|---|---
Frob-SAM | 5e-3 | 5e-3 | 5e-3 | 1e-2 | 5e-3 | 5e-3 | 1e-2 | 5e-3 | 1e-2
Table 5: Choice of $\rho$ for Frob-SAM for different datasets.

###  I.2 Additional Plots

All training plots are shown in Figure 3.

Figure 3: Training plots for all settings. One standard error is shaded.