# An SDE for Modeling SAM: Theory and Insights

Enea Monzio Compagnoni  Luca Biggio  Antonio Orvieto  Frank Norbert Proske
Hans Kersting  Aurelien Lucchi

###### Abstract

We study the SAM (Sharpness-Aware Minimization) optimizer which has recently
attracted a lot of interest due to its increased performance over more
classical variants of stochastic gradient descent. Our main contribution is
the derivation of continuous-time models (in the form of SDEs) for SAM and two
of its variants, both for the full-batch and mini-batch settings. We
demonstrate that these SDEs are rigorous approximations of the real discrete-
time algorithms (in a weak sense, scaling linearly with the learning rate).
Using these models, we then offer an explanation of why SAM prefers flat
minima over sharp ones – by showing that it minimizes an implicitly
regularized loss with a Hessian-dependent noise structure. Finally, we prove
that SAM is attracted to saddle points under some realistic conditions. Our
theoretical results are supported by detailed experiments.

Machine Learning, ICML

##  1 Introduction

Optimization plays a fundamental role in the performance of machine learning
models. The core problem it addresses is the minimization of the following
optimization problem:

$\min_{x\in{\mathbb{R}}^{d}}\left[f(x):=\frac{1}{N}\sum_{i=1}^{N}f_{i}(x)\right],$
(1)

where $f,f_{i}:{\mathbb{R}}^{d}\to{\mathbb{R}}$ for $i=1,\dots,N$. In machine
learning, $f$ is an empirical risk (or loss) function where $f_{i}$ are the
contributions due to the $i$-th data point. In this notation,
$x\in{\mathbb{R}}^{d}$ is a vector of trainable parameters and $N$ is the size
of the dataset.

Solving Eq. (1) is typically achieved via Gradient Descent (GD) methods that,
starting from a given estimate $x_{0}$, iteratively update an estimate $x_{k}$
as follows,

$x_{k+1}=x_{k}-\eta\nabla f(x_{k}),$ (2)

where $\eta>0$ is the learning rate. Since $\nabla f(x)$ requires computing
the average of the $N$ gradients $\nabla f_{i}(x)$ (which is computationally
expensive for large datasets where $N\gg 1$), it is common to instead replace
$\nabla f(x_{k})$ with a gradient estimated on a subset $\gamma_{k}$ of size
$B\geq 1$ of the dataset which is called a mini-batch. The resulting algorithm
is known as Stochastic Gradient Descent (SGD) whose update is

$x_{k+1}=x_{k}-\eta\nabla f_{\gamma_{k}}(x_{k}),$ (3)

where $\\{\gamma_{k}\\}$ are modelled here as i.i.d. random variables
uniformly distributed and taking value in $\\{1,\cdots,N\\}$.

Recently, Foret et al. (2021) proposed a stochastic optimizer known as
Sharpness-Aware Minimization (SAM), which yields significant performance gains
in various fields such as computer vision and natural language processing
(Bahri et al., 2022; Foret et al., 2021). The general idea behind SAM is to
seek parameters in low-loss regions that have a flatter curvature, which has
been shown to improve the generalization of the model (Hochreiter &
Schmidhuber, 1997; Keskar et al., 2017; Dziugaite & Roy, 2017; Jiang et al.,
2019). For a radius $\rho>0$, the iteration of SAM is

$x_{k+1}=x_{k}-\eta\nabla f_{\gamma_{k}}\left(x_{k}+\rho\frac{\nabla
f_{\gamma_{k}}(x_{k})}{\lVert\nabla f_{\gamma_{k}}(x_{k})\rVert}\right).$ (4)

Numerous works have studied SAM and proposed variants such as ESAM (Du et al.,
2022), ASAM (Kwon et al., 2021), GSAM (Zhuang et al., 2022), as well as Random
SAM and Variational SAM (Ujváry et al., 2022). Since SAM is hard to treat
theoretically, (Andriushchenko & Flammarion, 2022) introduced USAM which is
more easily analyzable as it drops the gradient normalization in Eq. (4), thus
yielding the following update:

$x_{k+1}=x_{k}-\eta\nabla f_{\gamma_{k}}\left(x_{k}+\rho\nabla
f_{\gamma_{k}}(x_{k})\right).$ (5)

Before analyzing the full version of SAM, we first take a smaller step toward
it by considering a variant with a deterministic normalization factor. We call
the resulting algorithm DNSAM (Deterministic Normalization SAM), whose update
step is

$x_{k+1}=x_{k}-\eta\nabla f_{\gamma_{k}}\left(x_{k}+\rho\frac{\nabla
f_{\gamma_{k}}(x_{k})}{\lVert\nabla f(x_{k})\rVert}\right).$ (6)

We will demonstrate both theoretically and empirically that DNSAM is a better
proxy of SAM than USAM. However, we do not claim that DNSAM is an algorithm to
be used in practice as its update requires the calculation of the full
gradient of the loss.

Following the theoretical framework of (Li et al., 2017), our work provides
the first _formal_ derivation of the SDEs of DNSAM, USAM, and SAM. Formally,
such continuous-time models are weak approximations (i.e. approximations in
distribution) of their respective discrete-time models. Importantly, SDE
models are not meant to be used as practical implementations since they have
to be discretized, giving rise to their discrete-time counterparts. Instead,
continuous-time models have typically been used in the recent literature to
derive novel insights about the discrete algorithms, see e.g. (Su et al.,
2014; Li et al., 2017).

We make the following contributions:

  1. 1.

Small $\rho$ regime. If $\rho=\mathcal{O}(\eta)$, we show that USAM, DNSAM,
and SAM essentially behave like SGD.

  2. 2.

Moderate $\rho$ regime. For $\rho=\mathcal{O}(\sqrt{\eta})$, we derive an SDE
model of USAM (7), of DNSAM (10), and of SAM (11). These can be interpreted as
the SDE of SGD on an implicitly regularized loss and with an additional
implicit curvature-induced noise. Leveraging these results, we demonstrate
that the additional noise is driven by the Hessian of the loss so that the
noise of the processes is larger in sharp minima. This is a key factor that
leads SAM and its variants to prefer flatter minima where the additional noise
decreases. However, while larger values of $\rho$ increase the noise of the
process, it also amplifies the implicit bias of the optimizer toward critical
points independently of whether they are minima, saddles, or maxima.

  3. 3.

Both in the full and mini-batch versions, USAM and SAM have very different
implicit regularizations.

  4. 4.

USAM might be attracted by saddles if $\rho$ is too large. Differently, for
any $\rho>0$, DNSAM and SAM might be attracted by saddles but eventually,
escape them after a long time. Thus, DNSAM is a more reliable model to
theoretically study SAM than USAM.

  5. 5.

Empirical validation. We empirically validate the proposed SDEs on several
models and landscapes commonly studied in the optimization and machine
learning communities.

In order to gain further insights from these continuous-time models, we also
study their behaviors on quadratic losses. The latter are commonly used to
model the landscape in the proximity of a critical point (Ge et al., 2015;
Levy, 2016; Jin et al., 2017; Poggio et al., 2017; Mandt et al., 2017b),
including several recent works that studied SAM (Bartlett et al., 2022; Wen et
al., 2023). This leads us to the following important observations:

  1. 1.

ODE - Pitfalls. After noticing that the ODE of SAM and DNSAM coincide, we
derive precise conditions under which SAM and USAM converge to the origin even
when it is a saddle or a maximum.

  2. 2.

SDE - Pitfalls. We derive the stationary distribution of the USAM SDE and find
that even this model is attracted by saddles under the same condition on
$\rho$ as found for the ODE 111Of course, the SDE does not point-wise converge
to the origin but rather oscillates around it with a certain variance.. In
contrast to USAM, we find that the dynamics of DNSAM is more complex: while a
certain region centered at the origin behaves like an attractor, the origin
itself repulses the dynamics away as the volatility rapidly increases to
infinity. This behavior of DNSAM is consistent with what was empirically
reported in (Kaddour et al., 2022) about SAM being able to get stuck around
saddles. To the best of our knowledge, this is the first time that these
potential pitfalls are formally demonstrated.

  3. 3.

Empirical validation. We empirically validate our claims for the quadratic
loss as well as other models.

##  2 Related Work

##### Theoretical Understanding of SAM

The current understanding of the dynamics of SAM and USAM is still limited.
Prior work includes the recent work by (Bartlett et al., 2022) that shows
that, for convex quadratics, SAM converges to a cycle oscillating between the
sides of the minimum in the direction with the largest curvature. For the non-
quadratic case, they also show that the dynamics drifts towards wider minima.
A concurrent work by (Wen et al., 2023) makes similar findings to (Bartlett et
al., 2022) as well as provides further insights regarding which notion of
sharpness SAM regularizes. Interestingly, the behavior of full-batch and mini-
batch SAM is intrinsically different. The former minimizes the largest
eigenvalue of the hessian of the loss, while the latter tries to uniformly
reduce the magnitude of the trace of the hessian of the loss. More
interestingly, (Wen et al., 2023) show how the dynamics of $1$-SAM can be
divided into two phases. The first phase follows the gradient flow with
respect to the loss until the dynamics approaches a manifold of minimizers. In
the second phase, the dynamics is driven towards parts of the landscape with a
lower trace of the hessian of the loss. (Rangwani et al., 2022) showed that
USAM could in some cases escape saddles faster than SGD. We however note that
their analysis is not completely formal as it relies on prior results by
(Daneshmand et al., 2018) which were specifically derived for SGD, not for
USAM. On our side, we here provide a more complete and rigorous description
that shows that USAM can be much slower than SGD at escaping a saddle.
Finally, a concurrent work (Kim et al., 2023) informally derived an SDE for
USAM around critical points, which relies on approximating the objective
function by a quadratic function. We remark that the authors did not formally
derive any guarantee showing the SDE closely approximates the true discrete-
time algorithm. In contrast, we formally and empirically demonstrate the
validity of our SDEs. In addition, our SDEs and analyses do not require the
quadratic approximation assumption made by (Kim et al., 2023) and are instead
valid for the entire trajectory of an optimizer, including, of course, around
critical points.

##### ODE Approximations

Continuous-time models in the form of (stochastic) differential equations are
a well-established tool to study discrete-time optimizers; see e.g. Helmke &
Moore (1994) and Kushner & Yin (2003). In machine learning, such models have
lately received increasing interest to study both deterministic and stochastic
optimizers. A notable reference is the work by Su et al. (2014) that derives a
second-order ODE to model Nesterov’s accelerated gradient. ODE models have
also been used recently to study SAM. This includes the work of Wen et al.
(2023, Section 4.2) discussed above, as well as Andriushchenko & Flammarion
(2022, Appendix B.1). Importantly we highlight two significant differences
with our work. First, our analysis focuses on the stochastic setting for which
we derive SDEs. Second, the ODE representations used in Wen et al. (2023) only
hold formally in the limit $\rho\rightarrow 0$, which is not the case in
practical settings where $\rho>0$. In contrast, our analysis allows for
significantly larger values of $\rho$, more precisely
$\rho=\mathcal{O}(\sqrt{\eta})$. Last but not least, neither of these papers
empirically validates the ODEs they derived.

##### SDE Approximations of Stochastic Optimizers.

For _stochastic_ optimizers, Li et al. (2017, 2019) derived an SDE that
provably approximates SGD (in the weak sense, i.e. in distribution). The
validity of this SDE model was experimentally tested in (Li et al., 2021).
Similar results are derived for ASGD by (An et al., 2020), and for RMSprop and
Adam by Malladi et al. (2022). In this paper, we derive an SDE approximation
for SAM, DNSAM, and USAM. The proof technique employed in our work (as well as
in An et al. (2020); Malladi et al. (2022)) relies on the theoretical
framework established by Li et al. (2017, 2019) (which itself requires
Assumption A.3 to hold). SDE approximations have also been derived for
different types of noise. This includes heavy-tailed noise that is shown to be
a good model for the noise of SGD in Simsekli et al. (2019), although the
evidence is still somewhat contested (Panigrahi et al., 2019; Xie et al.,
2021; Li et al., 2021). Zhou et al. (2020) also derived a Lévy-driven
stochastic differential equation to model the non-gaussianity of the noise,
which however does not enjoy any type of known theoretical approximation
guarantee. Finally, fractional Brownian noise, a generalization of Brownian
noise that allows for correlation, was used by (Lucchi et al., 2022).

##### Applications of SDE Approximations.

Continuous-time models are valuable analysis tools to study and design new
optimization methods. For instance, one concrete application of such models is
the use of _stochastic optimal control_ to select the learning rate (Li et
al., 2017, 2019) or the batch size (Zhao et al., 2022). In addition, _scaling
rules_ to adjust the optimization hyperparameters w.r.t. the batch size can be
recovered from SDE models (Malladi et al., 2022). Apart from these algorithmic
contributions, SDE approximation can be useful to better understand stochastic
optimization methods. In this regard, (Jastrzebski et al., 2018) analyzed the
factors influencing the minima found by SGD, and (Orvieto & Lucchi, 2019)
derived convergence bounds for mini-batch SGD and SVRG. (Smith et al., 2020)
used an SDE model to distinguish between ”noise-dominated” and ”curvature-
dominated” regimes of SGD. Yet another example is the study of _escape times_
of SGD from minima of different sharpness (Xie et al., 2021). Moreover, (Li et
al., 2020) and (Kunin et al., 2021) studied the _dynamics_ of the SDE
approximation under some symmetry assumptions. Finally, SDEs can be studied
through the lens of various tools in the field of stochastic calculus, e.g.
the Fokker–Planck equation gives access to the stationary distribution of a
stochastic process. Such tools are for instance valuable in the field of
Bayesian machine learning (Mandt et al., 2017a). For additional references,
see (Kushner & Yin, 2003; Ljung et al., 2012; Chen et al., 2015; Mandt et al.,
2015; Chaudhari & Soatto, 2018; Zhu et al., 2019; He et al., 2018; An et al.,
2020).

##  3 Formal Statements & Insights: The SDEs

In this section, we present the general formulations of the SDEs of USAM,
DNSAM, and SAM. Due to the technical nature of the analysis, we refer the
reader to the Appendix for the complete formal statements and proofs. For
didactic reasons, we provide simplified versions under mild additional
assumptions in the main paper.

######  Definition 3.1 (Weak Approximation).

Let $G$ denote the set of continuous functions
$\mathbb{R}^{d}\rightarrow\mathbb{R}$ of at most polynomial growth, i.e. $g\in
G$ if there exists positive integers $\kappa_{1},\kappa_{2}>0$ such that
$|g(x)|\leq\kappa_{1}\left(1+|x|^{2\kappa_{2}}\right)$, for all
$x\in\mathbb{R}^{d}$. Then, we say that a continuous-time stochastic process
$\left\\{X_{t}:t\in[0,T]\right\\}$ is an order $\alpha$ weak approximation of
a discrete stochastic process $\left\\{x_{k}:k=0,\ldots,N\right\\}$ if for
every $g\in G$, there exists a positive constant $C$, independent of $\eta$,
such that
$\max_{k=0,\ldots,N}\left|\mathbb{E}g\left(x_{k}\right)-\mathbb{E}g\left(X_{k\eta}\right)\right|\leq
C\eta^{\alpha}.$

This definition comes from the field of numerical analysis of SDEs, see
(Mil’shtein, 1986). Consider the case where $g(x)=\lVert x\rVert^{j}$, then
the bound limits the difference between the $j$-th moments of the discrete and
the continuous process.

In Theorem A.7 (USAM), Theorem A.11 (DNSAM), and Theorem A.16 (SAM), we prove
that if $\rho=\mathcal{O}(\eta)$ (small $\rho$ regime), the SDE of SGD (Eq.
(18)) is also an order 1 weak approximation for USAM, DNSAM, and SAM. In
contrast, in the more realistic moderate $\rho$ regime where
$\rho=\mathcal{O}(\sqrt{\eta})$, Eq. (18) is no longer an order 1 weak
approximation for any the models we analyze. Under such a condition, we
recover more insightful SDEs.

###  3.1 USAM SDE

######  Theorem 3.2 (USAM SDE - Informal Statement of Theorem A.4).

Under sufficient regularity conditions and $\rho=\mathcal{O}(\sqrt{\eta})$ the
solution of the following SDE is an order 1 weak approximation of the discrete
update of USAM (5):

$\displaystyle dX_{t}$
$\displaystyle=-\nabla\tilde{f}^{\text{USAM}}(X_{t})dt+$ (7)
$\displaystyle\sqrt{\eta\left(\Sigma^{\text{SGD}}(X_{t})+\rho\left(\tilde{\Sigma}(X_{t})+\tilde{\Sigma}(X_{t})^{\top}\right)\right)}dW_{t},$

where $\tilde{\Sigma}(x)$ is defined as

$\displaystyle\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\cdot\right.$ (8)
$\displaystyle\left.\left(\mathbb{E}\left[\nabla^{2}f_{\gamma}(x)\nabla
f_{\gamma}(x)\right]-\nabla^{2}f_{\gamma}(x)\nabla
f_{\gamma}(x)\right)^{\top}\right]$

and
$\tilde{f}^{\text{USAM}}(x):=f(x)+\frac{\rho}{2}\mathbb{E}\left[\lVert\nabla
f_{\gamma}(x)\rVert^{2}_{2}\right].$

To have a more direct comparison with the SDE of SGD (Eq. (18)), we prove
Corollary 3.3, a consequence of Theorem 3.2 that provides a more interpretable
SDE for USAM.

######  Corollary 3.3 (Informal Statement of Corollary A.6).

Under the assumptions of Theorem (3.2) and assuming a constant gradient noise
covariance, i.e. $\nabla f_{\gamma}(x)=\nabla f(x)+Z$ such that $Z$ is a noise
vector that does not depend on $x$, the solution of the following SDE is the
order 1 weak approximation of the discrete update of USAM (5):

$\displaystyle dX_{t}=$ $\displaystyle-\nabla\tilde{f}^{\text{USAM}}(X_{t})dt$
(9) $\displaystyle+$
$\displaystyle\left(I_{d}+\rho\nabla^{2}f(X_{t})\right)\left(\eta\Sigma^{\text{SGD}}\left(X_{t}\right)\right)^{1/2}dW_{t},$

where $\tilde{f}^{\text{USAM}}(x):=f(x)+\frac{\rho}{2}\lVert\nabla
f(x)\rVert^{2}_{2}.$

Corollary 3.3 shows that the dynamics of USAM is equivalent to that of SGD on
a regularized loss and with an additional noise component that depends on the
curvature of the landscape (captured by the term $\nabla^{2}f$).

###  3.2 DNSAM: A step towards SAM

######  Theorem 3.4 (DNSAM SDE - Informal Statement of Theorem A.9).

Under sufficient regularity conditions, assuming a constant gradient noise
covariance, i.e. $\nabla f_{\gamma}(x)=\nabla f(x)+Z$ such that $Z$ is a noise
vector that does not depend on $x$, and $\rho=\mathcal{O}(\sqrt{\eta})$ the
solution of the following SDE is the order 1 weak approximation of the
discrete update of DNSAM (6):

$\displaystyle dX_{t}=$
$\displaystyle-\nabla\tilde{f}^{\text{DNSAM}}(X_{t})dt$ (10) $\displaystyle+$
$\displaystyle\left(I_{d}+\rho\frac{\nabla^{2}f(X_{t})}{\lVert\nabla
f(X_{t})\rVert_{2}}\right)\left(\eta\Sigma^{\text{SGD}}(X_{t})\right)^{\frac{1}{2}}dW_{t}$

and $\tilde{f}^{\text{DNSAM}}(x)=f(x)+\rho\lVert\nabla f(x)\rVert_{2}$.

Theorem 3.4 shows that similarly to USAM, the dynamics of DNSAM is equivalent
to that of SGD on a regularized loss with an additional noise component that
depends on the curvature of the landscape. However, we notice that, unlike
USAM, the volatility component explodes near critical points.

###  3.3 SAM SDE

######  Theorem 3.5 (SAM SDE - Informal Statement of Theorem A.12).

Under sufficient regularity conditions and $\rho=\mathcal{O}(\sqrt{\eta})$ the
solution of the following SDE is the order 1 weak approximation of the
discrete update of SAM (4):

$\displaystyle dX_{t}=$ $\displaystyle-\nabla\tilde{f}^{\text{SAM}}(X_{t})dt$
(11) $\displaystyle+$
$\displaystyle\sqrt{\eta\left(\Sigma^{\text{SGD}}(X_{t})+\rho\left(\hat{\Sigma}(X_{t})+\hat{\Sigma}(X_{t})^{\top}\right)\right)}dW_{t}$

where $\hat{\Sigma}(x)$ is defined as

$\displaystyle\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\cdot\right.$ (12)
$\displaystyle\left.\left({\mathbb{E}}\left[\frac{\nabla^{2}f_{\gamma}(x)\nabla
f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert_{2}}\right]-\frac{\nabla^{2}f_{\gamma}(x)\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert_{2}}\right)^{\top}\right]$

and $\tilde{f}^{\text{SAM}}(x):=f(x)+\rho\mathbb{E}\left[\lVert\nabla
f_{\gamma}(x)\rVert_{2}\right].$

To have a more direct comparison with the SDE of SGD (Eq. (18)), we derive a
corollary of Theorem 3.5 that provides a more insightful SDE for SAM.

######  Corollary 3.6 (Informal Statement of Corollary A.15).

Under the assumptions of Theorem 3.5 and assuming a constant gradient noise
covariance, i.e. $\nabla f_{\gamma}(x)=\nabla f(x)+Z$ such that $Z$ is a noise
vector that does not depend on $x$, the solution of the following SDE is an
order 1 weak approximation of the discrete update of SAM (4)

$\displaystyle dX_{t}$ $\displaystyle=-\nabla\tilde{f}^{\text{SAM}}(X_{t})dt+$
(13) $\displaystyle\sqrt{\eta\left(\Sigma^{\text{SGD}}(X_{t})+\rho
H_{t}\left(\bar{\Sigma}(X_{t})+\bar{\Sigma}(X_{t})^{\top}\right)\right)}dW_{t}$

where $H_{t}:=\nabla^{2}f(X_{t})$ and $\bar{\Sigma}(x)$ is defined as

$\displaystyle\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\cdot\right.$ (14)
$\displaystyle\left.\left({\mathbb{E}}\left[\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert_{2}}\right]-\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert_{2}}\right)^{\top}\right],$

and $\tilde{f}^{\text{SAM}}(x):=f(x)+\rho\mathbb{E}\left[\lVert\nabla
f_{\gamma}(x)\rVert_{2}\right].$

We note that the regularization term of SAM is the expected norm of some
gradient. While one can of course use sampling in order to simulate the SDE in
Eq. (13), it results in an additional computational cost, which is worth
highlighting.

###  3.4 Comparison: USAM vs (DN)SAM

The analyses of the SDEs we derived (Eq. (9), Eq. (13), and Eq. (10)) reveal
that all three algorithms are implicitly minimizing a regularized loss with an
additional injection of noise (in addition to the SGD noise). While the
regularized loss is $\frac{\rho}{2}\lVert\nabla f(x)\rVert^{2}_{2}$ for USAM,
it is $\rho\lVert\nabla f(x)\rVert_{2}$ (not squared) for DNSAM, and
$\rho\mathbb{E}\left[\lVert\nabla f_{\gamma}(x)\rVert_{2}\right]$ for SAM.
Therefore, when the process is closer to a stationary point, the
regularization is stronger for (DN)SAM while it is the opposite when it is far
away.

Regarding the additional noise, we observe that it is curvature-dependent as
the Hessian appears in the expression of all volatility terms. Note that the
sharper the minimum, the larger the noise contribution from the Hessian. This
phenomenon is more extreme for DNSAM where the volatility is scaled by the
inverse of the norm of the gradient which drives the volatility to explode as
it approaches a critical point. Based on the SAM SDE, it is clear that the
diffusion term is more regular than that of DNSAM (in the sense that the
denominator does not vanish). Therefore, SAM is intrinsically less volatile
than DNSAM near a critical point. In contrast, we note that the SAM dynamics
exhibits more randomness than USAM which in turn is more noisy than SGD. These
theoretical insights are validated experimentally in Section 5. Therefore, it
is intuitive that SAM and its variants are more likely to stop or oscillate in
a flat basin and more likely to escape from sharp minima than SGD.

We conclude with a discussion of the role of $\rho$. On one hand, larger
values of $\rho$ increase the variance of the process. On the other hand, they
also increase the marginal importance of the factor
$\frac{\rho}{2}\lVert\nabla f(x)\rVert^{2}_{2}$ (USAM) and $\rho\lVert\nabla
f(x)\rVert_{2}$ (DNSAM), and $\rho\mathbb{E}\left[\lVert\nabla
f_{\gamma}(x)\rVert_{2}\right]$ (SAM), which for sufficiently large $\rho$
might overshadow the marginal relevance of minimizing $f$ and thus implicitly
bias the optimizer toward any point with zero gradients, including saddles and
maxima. We study this pitfall in detail for the quadratic case in the next
section and verify it experimentally in Section 5 for other models as well.
See Table 2 and Table 3 for a detailed summary.

##  4 Behavior Near Saddles - Theory

In this section, we leverage the ODEs (modeling the full-batch algorithms) and
SDEs (modeling the mini-batch algorithms) to study the behavior of SAM and its
variants near critical points. We especially focus on saddle points that have
been a subject of significant interest in machine learning (Jin et al., 2017,
2021; Daneshmand et al., 2018). We consider a quadratic loss (which as
mentioned earlier is a common model to study saddle points) of the form
$f(x)=\frac{1}{2}x^{\top}Hx$. W.l.o.g. we assume that the Hessian matrix $H$
is diagonal 222Recall that symmetric matrices can be diagonalized. and denote
the eigenvalues of $H$ by $(\lambda_{1},\dots,\lambda_{d})$ where
$\lambda_{1}\geq\lambda_{1}\geq\dots\geq\lambda_{d}$. If there are negative
eigenvalues, we denote by $\lambda_{*}$ the largest negative eigenvalue.

###  4.1 USAM ODE

We study the deterministic dynamics of USAM on a quadratic which is defined as

$dX_{t}=-H\left(I_{d}+\rho H\right)X_{t}dt\Rightarrow
X^{j}_{t}=X^{j}_{0}e^{-\lambda_{j}(1+\rho\lambda_{j})t}.$ (15)

Therefore, it is obvious (see Lemma C.1) that, if all the eigenvalues of $H$
are positive, for all $\rho>0$, we have that
$X^{j}_{t}\overset{t\rightarrow\infty}{\rightarrow}0,\quad\forall
j\in\\{1,\dots,d\\}$. In particular, we notice that, since
$e^{-\lambda_{j}(1+\rho\lambda_{j})t}<e^{-\lambda_{j}t}$, such convergence to
$0$ is faster for the flow of USAM than for the gradient flow. More
interestingly, if $\rho$ is too large, the following result states that the
deterministic dynamics of USAM might be attracted by a saddle or even a
maximum.

######  Lemma 4.1 (Informal Statement of Lemma C.2).

Let $H$ have at least one strictly negative eigenvalue. Then, for all
$\rho>-\frac{1}{\lambda_{*}},\quad
X^{j}_{t}\overset{t\rightarrow\infty}{\rightarrow}0,\quad\forall
j\in\\{1,\dots,d\\}$.

Therefore, if $\rho$ is not chosen appropriately, USAM might converge to
$0\in\mathbb{R}^{d}$, even if it is a saddle point or a maximum, which is very
undesirable. Of course, we observe that if $\rho<\frac{1}{\lambda_{*}}$, USAM
will diverge from the saddle (or maximum), which is desirable. Interestingly,
we also notice that since
$e^{-\lambda_{j}(1+\rho\lambda_{j})t}<e^{-\lambda_{j}t}$, USAM will escape the
saddle but more slowly than the gradient flow.

###  4.2 USAM SDE

Based on Corollary A.6, if we assume that
$\Sigma^{\text{SGD}}=\varsigma^{2}I_{d}$, the SDE of USAM on a quadratic is
given by

$dX_{t}=-H\left(I_{d}+\rho H\right)X_{t}dt+\left[(I_{d}+\rho
H)\sqrt{\eta}\varsigma\right]dW_{t}.$ (16)

######  Theorem 4.2 (Stationary distribution - Theorem C.3 and Theorem C.4).

If all the eigenvalues of $H$ are positive, i.e. $0$ is a minimum, we have
that for any $\rho>0$, $\forall i\in\\{1,\cdots,d\\}$, the stationary
distribution of Eq. (16) is

$\displaystyle
P\left(x\right)=\sqrt{\frac{\lambda_{i}}{\pi\eta\varsigma^{2}}\frac{1}{1+\rho\lambda_{i}}}\exp\left[-\frac{\lambda_{i}}{\eta\varsigma^{2}}\frac{1}{1+\rho\lambda_{i}}x^{2}\right].$
(17)

If there exists a negative eigenvalue, this formula does not, in general,
parametrize a probability distribution. However, if
$\rho>-\frac{1}{\lambda_{*}}$, Eq. (17) is still the stationary distribution
of Eq.(16), $\forall i\in\\{1,\cdots,d\\}$.

Theorem 4.2 states that in case the origin is a saddle (or a maximum) and
$\rho$ is small enough, the stationary distribution of USAM is divergent at
infinity, meaning that the process will escape the bad critical point, which
is desirable. In such a case, the escape from the saddle is however slower
than SGD as the variance in the direction of negative eigenvalues, e.g. the
escape directions, is smaller. However, if $\rho$ is too large, then the
dynamics of the USAM SDE will oscillate around the origin even if this is a
saddle or a maximum, which is undesirable. This is consistent with the results
derived for the SDE of USAM in Section 4.1. There, we found that under the
very same condition on $\rho$, the USAM ODE converges to $0$ even when it is a
saddle or a maximum.

###  4.3 SAM ODE

We now provide insights on the dynamics of the SAM ODE on a quadratic with
Hessian $H$.

######  Lemma 4.3 (Lemma C.6).

For all $\rho>0$, if $H$ is PSD (Positive Semi-Definite), the origin is
(locally) asymptotically stable. Additionally, if $H$ is not PSD and $\lVert
HX_{t}\rVert\leq-\rho\lambda_{*}$, then the origin is still (locally)
asymptotically stable.

Lemma 4.3 demonstrates that USAM and SAM have completely different behaviors.
For USAM, Lemma 4.1 shows that selecting $\rho$ small enough would prevent the
convergence towards a saddle or a maximum. In contrast, Lemma 4.3 shows that
for any value of $\rho$, if the dynamics of SAM is close enough to any
critical point, i.e. enters an attractor, it is attracted by it. We also
observe that if $\rho\rightarrow 0$, this attractor reduces to a single point,
i.e. the critical point itself.

To the best of our knowledge, this is the first time that these phenomena are
formally demonstrated. Importantly, these theoretical insights are consistent
with the experimental results of (Kaddour et al., 2022) that show how SAM
might get stuck around saddles.

Finally, by comparing the ODE of USAM (Eq. (15)) with that of SAM (Eq. (149)),
we observe that the dynamics of SAM is equivalent to that of USAM where the
radius $\rho$ has been scaled by $\frac{1}{\lVert HX_{t}\rVert}$. In a way,
while USAM has a fixed radius $\rho$, SAM has an time-dependent radius
$\frac{\rho}{\lVert HX_{t}\rVert}$ which is smaller than $\rho$ if the
dynamics is far from the origin ($\lVert HX_{t}\rVert>1$) and larger when it
is close to it ($\lVert HX_{t}\rVert<1$). Therefore, SAM converges to the
origin slower than USAM when it is far from it and it becomes faster as it
gets closer.

###  4.4 (DN)SAM SDE

We now provide insights on the dynamics of the DNSAM SDE on a quadratic with
Hessian $H$.

######  Observation 4.4 (Details in C.7).

We observe that for all $\rho>0$, there exists an $\epsilon>0$ such that if
$\lVert HX_{t}\rVert\in\left(\epsilon,-\rho\lambda_{*}\right)$, the dynamics
of $X_{t}$ is attracted towards the origin. If the eigenvalues are all
positive, the condition becomes $\lVert
HX_{t}\rVert\in\left(\epsilon,\infty\right)$. On the contrary, if $\lVert
HX_{t}\rVert<\epsilon$, then the dynamics is pushed away from the origin.

This insight suggests that if DNSAM is initialized close enough to a quadratic
saddle, it is attracted toward it, but is also repulsed by it if it gets too
close. This is due to the explosion of the volatility next to the origin. We
expect that this observation extends to SAM as well, and it remains to be
shown theoretically in future work. In the next section, we experimentally
verify that the dynamics gets cyclically pulled to $0$ and pushed away from
it, not only for the quadratic saddle but also for that of other models.

##  5 Experiments

The main goal of this experimental section is two-fold: 1) to verify the
validity of the theorems derived in Section 3, and 2) to validate the claims
made about the behavior of SAM and its variants near saddle points. The latter
requires us to use models, for which saddle points are known to be present
(Safran & Shamir, 2018), including for instance linear autoencoders (Kunin et
al., 2019).

###  5.1 Empirical Validation of the SDEs

We first experimentally validate the results of Corollary 3.3, Corollary 3.6,
and Theorem 3.4. To do so, we use two different test functions ($g(x)$ in Def.
(3.1)), which are $g_{1}(x):=\lVert x\rVert+\lVert\nabla f(x)\rVert$ and
$g_{2}(x):=f(x)$. We test on four models. The first model is a convex
quadratic landscape. The second task is a classification one on the Iris
Database (Dua & Graff, 2017) using a linear MLP with $1$ hidden layer. The
third is a classification task on the Breast Cancer Database (Dua & Graff,
2017) using a nonlinear MLP with $1$ hidden layer. The fourth is a Teacher-
Student model where the Teacher is a linear MLP with $20$ hidden layers and
the Student is a nonlinear MLP with $20$ hidden layers. Figure 2 uses the
first metric $g_{1}(x)$ to measure the maximum absolute error (across the
whole trajectory) of the SDEs of SGD, USAM, DNSAM, and SAM in approximating
the respective discrete algorithms. Additionally, we plot the same error if we
were to use the SDE of SGD to model/approximate the discrete iterates of SAM
and its variants. We observe that when $\rho=\eta$, the absolute error is
small in all cases, meaning that all the discrete iterates and SDEs behave
essentially in the same way. This supports our claim that if
$\rho=\mathcal{O}(\eta)$, the SDE of SGD is a good model for USAM, DNSAM, and
SAM (Theorem A.7, Theorem A.11, and Theorem A.16). When $\rho=\sqrt{\eta}$, we
see that the derived SDEs correctly approximate the respective discrete
algorithms, while the SDE of SGD has a significantly larger relative error,
which validates the results of Corollary 3.3, Theorem 3.4, and Corollary 3.6.
Although we do not have any theoretical guarantee for larger $\rho$, we
observe empirically that the modeling error is still rather low. Finally,
Figure 3 shows the evolution of the metric $g_{2}(x):=f(x)$ for the different
algorithms. We notice that all the SDEs are matching the respective
algorithms. In Appendix D.1.1, we provide evidence that failing to include the
correct diffusion terms in the USAM SDE Eq. (7) and the DNSAM SDE Eq. (10)
leads to less accurate models.

Figure 1: Comparison in terms of $g_{1}(x)$ with respect to $\rho$.

Finally, Figure 1 shows that, in the Quadratic case, DNSAM results in a much
closer approximation to SAM than other SDEs. Based on this observation and the
analyses of Section 4, we conclude that DNSAM is a better proxy to
theoretically study SAM than USAM. It however remains not advised to employ
DNSAM as a practical algorithm since its update rule requires the calculation
of the full gradient, see Eq. (6).

##### Interplay between noise, curvature, $\rho$, and suboptimality

Next, we check how the parameter $\rho$ and the curvature (measured by the
trace operator of the Hessian) influence the noise of the stochastic process
and its suboptimality. These insights substantiate the intuition that SAM and
its variants are more likely to escape sharp minima faster than SGD.

First of all, we fix the value of $\rho$ as well as a diagonal Hessian $H$
with random positive eigenvalues. Then, we study the loss for SGD, USAM,
DNSAM, and SAM as they optimize a convex quadratic loss of increasingly larger
curvature (i.e. larger Hessian magnitude). We observe that DNSAM exhibits a
loss that is orders of magnitude larger than that of SGD, with more variance,
and even more so as the curvature increases. Note that SAM behaves similarly
to DNSAM, but with less variance. Finally, USAM exhibits a similar pattern but
less pronounced. All the observations are consistent with the insights gained
from the covariance matrices in the respective SDEs. For details, we refer the
reader to Figure 8, Figure 9, and Figure 10 in Appendix.

In a similar experiment, we fix the Hessian as above and study the loss as we
increase $\rho$. Once again, we observe that DNSAM exhibits a larger loss with
more variance, and this is more and more clear as $\rho$ increases.
Observations similar to the above ones can be done for SAM and USAM. For
details, we refer the reader to Figure 11, 12 and Figure 13 in Appendix.

Finally, we note that SAM and its variant exhibit an additional implicit
curvature-induced noise compared to SGD. This leads to increased suboptimality
as well as a higher likelihood to escape sharp minima. We provide an
additional justification for this phenomenon in Observation C.5.

(a)  (b)

(c)  (d)

Figure 2: Comparison in terms of $g_{1}(x)$ with respect to $\rho$ \-
Quadratic (top left); Teacher-Student (top right); Deep linear class (lower
left); Deep Nonlinear class (lower right).

(a)  (b)

(c)  (d)

Figure 3: Comparison in terms of $g_{2}(x)$ with respect to time - Quadratic
(top left); Teacher-Student (top right); Deep linear class (lower left); Deep
Nonlinear class (lower right).

###  5.2 Behavior Near Saddles

In this section, we study the behavior of SAM and USAM (full batch versions),
and of PSAM, DNSAM, and PUSAM (perturbed gradient versions) near saddle
points. See Table 1 for more details.

##### Quadratic Landscape

We first empirically verify the insight gained in Section 4.4 — the dynamics
of DNSAM is attracted to the origin, but if it gets too close, it gets
repulsed away. For a quadratic saddle, in Figure 4 we show the distribution of
$10^{5}$ trajectories after $5\cdot 10^{4}$ iterations. These are distributed
symmetrically around the origin but the concentration is lower close to it.
While this is intuitive for the convex case (see Figure 14 in Appendix), it is
surprising for the saddle case: our insights are fully verified. The second
and third images of Figure 4 show that all the trajectories are initialized
outside of a certain ball around the origin and then they get pulled inside
it. Then, we see that the number of points outside this ball increased again
and the right-most image shows the number of points jumping in and out of it.
This shows that there is a cyclical dynamics towards and away from the origin.
Of course, all the points eventually escape the saddle, but much more slowly
than what would happen under the dynamics of SGD where the trajectories would
not even get close to the origin in the first place. In Figure 15 in Appendix,
we show the behavior of several optimizers when initialized in an escaping
direction from the saddle and we observe that full-batch SAM is attracted by
the saddle while the others are able to escape it. Interestingly, PSAM is
anyway slower than SGD in escaping. Figure 15 in Appendix shows that full-
batch SAM and PSAM cannot escape the saddle if it is too close to it, while
DNSAM can if it is close enough to enjoy a spike in volatility. More details
are in the Appendix D.2.

(a)  (b)

(c)  (d)

Figure 4: Quadratic Saddle - Top Left: Distribution points around the origin
are scarcer near to the origin; Top Right: Number of trajectories outside a
small ball around the origin increases over time; Lower Left: All the
trajectories eventually enter the ball and then start exiting it; Lower Right:
There is a constant oscillation of points in and out of the ball.

##### Linear Autoencoder

Inspired by the insights gained so far, we study the behavior of SAM when it
is initialized close to the saddle present at the origin of the linear
autoencoder introduced by (Kunin et al., 2019). The top-left of Figure 5 shows
the evolution of the loss as we optimize it with SAM starting from different
starting points closer and closer to the saddle in the origin. The scalar
$\sigma$ parametrizes how close the initialization is to the origin. We
observe that when SAM starts sufficiently far from it ($\sigma\geq 0.005$), it
optimizes immediately, while the closer it is initialized to it, the more it
stays around it, up to not being able to move at all ($\sigma\leq 0.001$).
Regarding DNSAM, in the top-right figure, we observe the same behavior, apart
from one case: if it is initialized sufficiently close to the origin, instead
of getting stuck there, it jumps away following a spike in volatility.
Differently, PSAM behaves more like SAM and is slower in escaping if $\sigma$
is lower. The bottom-right of Figure 5 shows the comparison with other
optimizers: SAM does not optimize the loss while the other optimizers do.
These findings are consistent with those observed in Figure 15 in Appendix for
the quadratic landscape. In Figure 16 in Appendix, we show a similar result
for a saddle landscape studied in (Lucchi et al., 2022). More details are in
Appendix D.3 and Appendix D.4, respectively. In both these experiments, we
observe the suboptimality patterns forecasted by our theory.

(a)  (b)

(c)  (d)

Figure 5: Autoencoder - Top-Left: SAM does not escape the saddle if it is too
close to it. Top-Right: DNSAM escapes if it is extremely close to the origin
thanks to a volatility spike. Bottom-Left: Like SAM, PSAM does not escape if
too close to the origin. Bottom-Right: DNSAM is the fastest to escape, while
SAM is stuck.

##  6 Discussion

###  6.1 Future Work

Inspired by (Malladi et al., 2022), it would be interesting to study possible
scaling rules for SAM and its variants, thus shedding light on the interplay
of the learning rate $\eta$, the batch size $B$, and $\rho$. Another direction
could be to use our SDEs to study the role of $\rho$ in balancing the
optimization speed and generalization properties of SAM. The insights gained
could be useful in improving SAM in terms of optimization speed and
generalization. Finally, we expect that the application of more analytical
tools to the SDEs of SAM and DNSAM will lead to further insights into SAM. It
would for instance be of particular interest to revisit claims made about
other optimizers via their SDE models (see “Applications of SDE
approximations” in Section 2). Hopefully, this will help to demystify the high
performance of SAM on large-scale ML problems.

###  6.2 Limitations

We highlight that modeling discrete-time algorithms via SDEs relies on
Assumption A.3. Furthermore, this setup cannot fully capture the regime of
large learning rates. As observed in (Li et al., 2021), a large $\eta$ or the
lack of certain conditions on $\nabla f$ and on the noise covariance matrix
might lead to an approximation failure. However, the authors claim that this
failure could be avoided by increasing the order of the weak approximation.
Additionally, most of our discussions are focused on the case where
$\rho=\mathcal{O}(\sqrt{\eta})$, which is not the only interesting setup, as
some authors use $\rho<\eta$. Finally, since our work is more theoretical in
nature, we did not aim at conducting SOTA experiments but rather focused on
improving the understanding of the dynamics of SAM. Thus, we analyzed
relatively simple models and landscapes that are relevant to the optimization
and machine learning community.

###  6.3 Conclusion

We proposed new continuous-time models (in the form of SDEs) for the SAM
optimizer and two variants. While the USAM variant was introduced in prior
work (Andriushchenko & Flammarion, 2022), the DNSAM variant we introduce is a
step between USAM and SAM, allowing us to gain further insights into the role
of the normalization. We formally proved (and experimentally verified) that
these SDEs approximate their real discrete-time counterparts; see Theorems
3.2–3.6 for the theory and Section 5.1 for the experiments. An interesting
side aspect of our analysis is that DNSAM appears to be a better surrogate
model to describe the dynamics of SAM than USAM: SAM and DNSAM share common
behaviors around saddles, they have more similar noise structures, and
experiments support these claims. Of course, by no means does this paper
intend to propose DNSAM as a new practical optimizer: it is instead meant to
be used for theoretical analyses.

The SDEs we derived explicitly decompose the learning dynamics (in the
parameter space) into a deterministic drift and a stochastic diffusion
coefficient which in itself reveals some novel insights: The drift coefficient
– by the definition of $\tilde{f}$ in Theorems 3.2–3.6 – exposes how the
ascent parameter $\rho$ impacts the average dynamics of SAM and its variants.
The diffusion coefficient, on the other hand, increases with the Hessian of
the loss – thereby implying that SAM and its variants are noisier in sharp
minima. This could be interpreted as an implicit bias towards flat minima (as
sharp minima will be more unstable due to the noise).

The continuous-time SDE models allow the application of tools from stochastic
calculus (e.g. integration and differentiation) to study the behavior of SAM.
As a start in this direction, we proved that the flow of USAM gets stuck
around saddles if $\rho$ is too large. In contrast, SAM oscillates around
saddles if initialized close to them but eventually slowly escapes them thanks
to the additional noise. Importantly, our claims are substantiated by
experiments on several models and invite further investigation to prevent a
costly waste of computation budget near saddle points.

##### Acknowledgement

We would like to thank the reviewers for their feedback which greatly helped
us improve this manuscript. Frank Proske acknowledges the financial support of
the Norwegian Research Council (project number 274410). Enea Monzio Compagnoni
and Aurelien Lucchi acknowledge the financial support of the Swiss National
Foundation, SNF grant No 207392. Hans Kersting thanks the European Research
Council for support through the ERC grant 724063.

## References

  * An et al. (2020) An, J., Lu, J., and Ying, L.  Stochastic modified equations for the asynchronous stochastic gradient descent.  _Information and Inference: A Journal of the IMA_ , 9(4):851–873, 2020.
  * Andriushchenko & Flammarion (2022) Andriushchenko, M. and Flammarion, N.  Towards understanding sharpness-aware minimization.  In _International Conference on Machine Learning_ , pp. 639–668. PMLR, 2022.
  * Bahri et al. (2022) Bahri, D., Mobahi, H., and Tay, Y.  Sharpness-aware minimization improves language model generalization.  _ACL 2022_ , 2022.
  * Bartlett et al. (2022) Bartlett, P. L., Long, P. M., and Bousquet, O.  The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima.  _arXiv preprint arXiv:2210.01513_ , 2022.
  * Chaudhari & Soatto (2018) Chaudhari, P. and Soatto, S.  Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks.  In _2018 Information Theory and Applications Workshop (ITA)_ , pp. 1–10. IEEE, 2018.
  * Chen et al. (2015) Chen, C., Ding, N., and Carin, L.  On the convergence of stochastic gradient mcmc algorithms with high-order integrators.  _Advances in neural information processing systems_ , 28, 2015.
  * Daneshmand et al. (2018) Daneshmand, H., Kohler, J., Lucchi, A., and Hofmann, T.  Escaping saddles with stochastic gradients.  In _International Conference on Machine Learning_ , pp. 1155–1164. PMLR, 2018.
  * Du et al. (2022) Du, J., Yan, H., Feng, J., Zhou, J. T., Zhen, L., Goh, R. S. M., and Tan, V. Y.  Efficient sharpness-aware minimization for improved training of neural networks.  _ICLR 2022_ , 2022.
  * Dua & Graff (2017) Dua, D. and Graff, C.  UCI machine learning repository, 2017.  URL http://archive.ics.uci.edu/ml.
  * Dziugaite & Roy (2017) Dziugaite, G. K. and Roy, D. M.  Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data.  In _Uncertainty in Artificial Intelligence_ , 2017.
  * Folland (2005) Folland, G. B.  Higher-order derivatives and taylor’s formula in several variables.  _Preprint_ , pp. 1–4, 2005.
  * Foret et al. (2021) Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B.  Sharpness-aware minimization for efficiently improving generalization.  _ICLR 2021_ , 2021.
  * Gardiner et al. (1985) Gardiner, C. W. et al.  _Handbook of stochastic methods_ , volume 3.  springer Berlin, 1985.
  * Ge et al. (2015) Ge, R., Huang, F., Jin, C., and Yuan, Y.  Escaping from saddle points—online stochastic gradient for tensor decomposition.  In _Conference on Learning Theory_ , pp. 797–842, 2015.
  * Gyöngy & Martínez (2001) Gyöngy, I. and Martínez, T.  On stochastic differential equations with locally unbounded drift.  _Czechoslovak Mathematical Journal, No. 4, p. 763–783, Vol. 51_ , 2001.
  * He et al. (2018) He, L., Meng, Q., Chen, W., Ma, Z.-M., and Liu, T.-Y.  Differential equations for modeling asynchronous algorithms.  In _Proceedings of the 27th International Joint Conference on Artificial Intelligence_ , IJCAI’18, pp. 2220–2226. AAAI Press, 2018.
  * Helmke & Moore (1994) Helmke, U. and Moore, J. B.  _Optimization and Dynamical Systems_.  Springer London, 1st edition, 1994.
  * Hochreiter & Schmidhuber (1997) Hochreiter, S. and Schmidhuber, J.  Flat minima.  _Neural computation_ , 9(1):1–42, 1997.
  * Jastrzebski et al. (2018) Jastrzebski, S., Kenton, Z., Arpit, D., Ballas, N., Fischer, A., Bengio, Y., and Storkey, A.  Three factors influencing minima in sgd.  _ICANN 2018_ , 2018.
  * Jiang et al. (2019) Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S.  Fantastic generalization measures and where to find them.  In _International Conference on Learning Representations_ , 2019.
  * Jin et al. (2017) Jin, C., Ge, R., Netrapalli, P., Kakade, S. M., and Jordan, M. I.  How to escape saddle points efficiently.  In _International Conference on Machine Learning_ , pp. 1724–1732. PMLR, 2017.
  * Jin et al. (2021) Jin, C., Netrapalli, P., Ge, R., Kakade, S. M., and Jordan, M. I.  On nonconvex optimization for machine learning: Gradients, stochasticity, and saddle points.  _J. ACM_ , 68(2), 2021.
  * Kaddour et al. (2022) Kaddour, J., Liu, L., Silva, R., and Kusner, M. J.  When do flat minima optimizers work?  _Advances in Neural Information Processing Systems_ , 35:16577–16595, 2022.
  * Keskar et al. (2017) Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T. P.  On large-batch training for deep learning: Generalization gap and sharp minima.  _ICLR 2017_ , 2017.
  * Kim et al. (2023) Kim, H., Park, J., Choi, Y., and Lee, J.  Stability analysis of sharpness-aware minimization.  _arXiv preprint arXiv:2301.06308_ , 2023.
  * Kunin et al. (2019) Kunin, D., Bloom, J., Goeva, A., and Seed, C.  Loss landscapes of regularized linear autoencoders.  In _International Conference on Machine Learning_ , pp. 3560–3569. PMLR, 2019.
  * Kunin et al. (2021) Kunin, D., Sagastuy-Brena, J., Ganguli, S., Yamins, D. L., and Tanaka, H.  Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics.  In _International Conference on Learning Representations_ , 2021.
  * Kushner & Yin (2003) Kushner, H. and Yin, G. G.  _Stochastic approximation and recursive algorithms and applications_ , volume 35.  Springer Science & Business Media, 2003.
  * Kwon et al. (2021) Kwon, J., Kim, J., Park, H., and Choi, I. K.  Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks.  In _International Conference on Machine Learning_ , pp. 5905–5914. PMLR, 2021.
  * Levy (2016) Levy, K. Y.  The power of normalization: Faster evasion of saddle points.  _arXiv preprint arXiv:1611.04831_ , 2016.
  * Li et al. (2017) Li, Q., Tai, C., and Weinan, E.  Stochastic modified equations and adaptive stochastic gradient algorithms.  In _International Conference on Machine Learning_ , pp. 2101–2110. PMLR, 2017.
  * Li et al. (2019) Li, Q., Tai, C., and Weinan, E.  Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations.  _The Journal of Machine Learning Research_ , 20(1):1474–1520, 2019.
  * Li et al. (2020) Li, Z., Lyu, K., and Arora, S.  Reconciling modern deep learning with traditional optimization analyses: the intrinsic learning rate.  In _Advances in Neural Information Processing Systems_ , 2020.
  * Li et al. (2021) Li, Z., Malladi, S., and Arora, S.  On the validity of modeling SGD with stochastic differential equations (SDEs).  In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), _Advances in Neural Information Processing Systems_ , 2021.
  * Ljung et al. (2012) Ljung, L., Pflug, G., and Walk, H.  _Stochastic approximation and optimization of random systems_ , volume 17.  Birkhäuser, 2012.
  * Lucchi et al. (2022) Lucchi, A., Proske, F., Orvieto, A., Bach, F., and Kersting, H.  On the theoretical properties of noise correlation in stochastic optimization.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Malladi et al. (2022) Malladi, S., Lyu, K., Panigrahi, A., and Arora, S.  On the SDEs and scaling rules for adaptive gradient algorithms.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Mandt et al. (2015) Mandt, S., Hoffman, M. D., Blei, D. M., et al.  Continuous-time limit of stochastic gradient descent revisited.  _NIPS-2015_ , 2015.
  * Mandt et al. (2017a) Mandt, S., Hoffman, M. D., and Blei, D. M.  Stochastic gradient descent as approximate Bayesian inference.  _J. Mach. Learn. Res._ , 18(1):4873–4907, 2017a.  ISSN 1532-4435.
  * Mandt et al. (2017b) Mandt, S., Hoffman, M. D., and Blei, D. M.  Stochastic gradient descent as approximate bayesian inference.  _JMLR 2017_ , 2017b.
  * Mao (2007) Mao, X.  _Stochastic differential equations and applications_.  Elsevier, 2007.
  * Mil’shtein (1986) Mil’shtein, G.  Weak approximation of solutions of systems of stochastic differential equations.  _Theory of Probability & Its Applications_, 30(4):750–766, 1986.
  * Orvieto & Lucchi (2019) Orvieto, A. and Lucchi, A.  Continuous-time models for stochastic optimization algorithms.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Panigrahi et al. (2019) Panigrahi, A., Somani, R., Goyal, N., and Netrapalli, P.  Non-Gaussianity of stochastic gradient noise.  _SEDL workshop at NeurIPS 2019_ , 2019.
  * Poggio et al. (2017) Poggio, T., Kawaguchi, K., Liao, Q., Miranda, B., Rosasco, L., Boix, X., Hidary, J., and Mhaskar, H.  Theory of deep learning iii: explaining the non-overfitting puzzle.  _arXiv preprint arXiv:1801.00173_ , 2017.
  * Rangwani et al. (2022) Rangwani, H., Aithal, S. K., Mishra, M., and Babu, R. V.  Escaping saddle points for effective generalization on class-imbalanced data.  _NeurIPS 2022_ , 2022.
  * Risken (1996) Risken, H.  Fokker-planck equation.  In _The Fokker-Planck Equation_ , pp. 63–95. Springer, 1996.
  * Safran & Shamir (2018) Safran, I. and Shamir, O.  Spurious local minima are common in two-layer relu neural networks.  In _International Conference on Machine Learning_ , pp. 4433–4441. PMLR, 2018.
  * Sagun et al. (2018) Sagun, L., Evci, U., Guney, V. U., Dauphin, Y., and Bottou, L.  Empirical analysis of the hessian of over-parametrized neural networks.  _ICLR 2018 Workshop Track_ , 2018.  URL https://openreview.net/forum?id=rJrTwxbCb.
  * Simsekli et al. (2019) Simsekli, U., Sagun, L., and Gurbuzbalaban, M.  A tail-index analysis of stochastic gradient noise in deep neural networks.  In _International Conference on Machine Learning_ , 2019.
  * Smith et al. (2020) Smith, S., Elsen, E., and De, S.  On the generalization benefit of noise in stochastic gradient descent.  In _International Conference on Machine Learning_ , 2020.
  * Su et al. (2014) Su, W., Boyd, S., and Candes, E.  A differential equation for modeling Nesterov’s accelerated gradient method: Theory and insights.  In _Advances in Neural Information Processing Systems_ , 2014.
  * Ujváry et al. (2022) Ujváry, S., Telek, Z., Kerekes, A., Mészáros, A., and Huszár, F.  Rethinking sharpness-aware minimization as variational inference.  In _OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop)_ , 2022.
  * Wen et al. (2023) Wen, K., Ma, T., and Li, Z.  How does sharpness-aware minimization minimize sharpness?  _ICLR 2023_ , 2023.
  * Xie et al. (2021) Xie, Z., Sato, I., and Sugiyama, M.  A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima.  In _International Conference on Learning Representations_ , 2021.
  * Zhao et al. (2022) Zhao, J., Lucchi, A., Proske, F. N., Orvieto, A., and Kersting, H.  Batch size selection by stochastic optimal control.  In _Has it Trained Yet? NeurIPS 2022 Workshop_ , 2022.
  * Zhou et al. (2020) Zhou, P., Feng, J., Ma, C., Xiong, C., Hoi, S. C. H., et al.  Towards theoretically understanding why sgd generalizes better than adam in deep learning.  _Advances in Neural Information Processing Systems_ , 33:21285–21296, 2020.
  * Zhu et al. (2019) Zhu, Z., Wu, J., Yu, B., Wu, L., and Ma, J.  The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects.  _ICML 2019_ , 2019.
  * Zhuang et al. (2022) Zhuang, J., Gong, B., Yuan, L., Cui, Y., Adam, H., Dvornek, N., Tatikonda, S., Duncan, J., and Liu, T.  Surrogate gap minimization improves sharpness-aware training.  _ICML 2022_ , 2022.

##  Appendix A Theoretical Framework - SDEs

In the subsequent proofs, we will make repeated use of Taylor expansions in
powers of $\eta$. To simplify the presentation, we introduce the shorthand
that whenever we write $\mathcal{O}\left(\eta^{\alpha}\right)$, we mean that
there exists a function $K(x)\in G$ such that the error terms are bounded by
$K(x)\eta^{\alpha}$. For example, we write

$b(x+\eta)=b_{0}(x)+\eta b_{1}(x)+\mathcal{O}\left(\eta^{2}\right)$

to mean: there exists $K\in G$ such that

$\left|b(x+\eta)-b_{0}(x)-\eta b_{1}(x)\right|\leq K(x)\eta^{2}.$

Additionally, let us introduce some notation:

  * •

A multi-index is $\alpha=\left(\alpha_{1},\alpha_{2},\ldots,\alpha_{n}\right)$
such that $\alpha_{j}\in\\{0,1,2,\ldots\\}$

  * •

$|\alpha|:=\alpha_{1}+\alpha_{2}+\cdots+\alpha_{n}$

  * •

$\alpha!:=\alpha_{1}!\alpha_{2}!\cdots\alpha_{n}!$

  * •

For $x=\left(x_{1},x_{2},\ldots,x_{n}\right)\in\mathbb{R}^{n}$, we define
$x^{\alpha}:=x_{1}^{\alpha_{1}}x_{2}^{\alpha_{2}}\cdots x_{n}^{\alpha_{n}}$

  * •

For a multi-index $\beta$,
$\partial_{\beta}^{|\beta|}f(x):=\frac{\partial^{|\beta|}}{\partial^{\beta_{1}}_{x_{1}}\partial^{\beta_{2}}_{x_{2}}\cdots\partial^{\beta_{n}}_{x_{n}}}f(x)$

  * •

We also denote the partial derivative with respect to $x_{i}$ by
$\partial_{e_{i}}$.

######  Lemma A.1 (Lemma 1 (Li et al., 2017)).

Let $0<\eta<1$. Consider a stochastic process $X_{t},t\geq 0$ satisfying the
SDE
$dX_{t}=b\left(X_{t}\right)dt+\eta^{\frac{1}{2}}\sigma\left(X_{t}\right)dW_{t}$
with $X_{0}=x\in\mathbb{R}^{d}$ and $b,\sigma$ together with their derivatives
belong to $G$. Define the one-step difference $\Delta=X_{\eta}-x$, then we
have 1.
$\mathbb{E}\Delta_{i}=b_{i}\eta+\frac{1}{2}\left[\sum_{j=1}^{d}b_{j}\partial_{e_{j}}b_{i}\right]\eta^{2}+\mathcal{O}\left(\eta^{3}\right)\quad\forall
i=1,\ldots,d$; 2.
$\mathbb{E}\Delta_{i}\Delta_{j}=\left[b_{i}b_{j}+\sigma\sigma_{(ij)}^{T}\right]\eta^{2}+\mathcal{O}\left(\eta^{3}\right)\quad\forall
i,j=1,\ldots,d$; 3.
$\mathbb{E}\prod_{j=1}^{s}\Delta_{\left(i_{j}\right)}=\mathcal{O}\left(\eta^{3}\right)$
for all $s\geq 3,i_{j}=1,\ldots,d$. All functions above are evaluated at $x$.

######  Theorem A.2 (Theorem 2 and Lemma 5, (Mil’shtein, 1986)).

Let the assumptions in Theorem A.4 hold and let us define
$\bar{\Delta}=x_{1}-x$ to be the increment in the discrete-time algorithm. If
in addition there exists $K_{1},K_{2},K_{3},K_{4}\in G$ so that 1.
$\left|\mathbb{E}\Delta_{i}-\mathbb{E}\bar{\Delta}_{i}\right|\leq
K_{1}(x)\eta^{2},\quad\forall i=1,\ldots,d$; 2.
$\left|\mathbb{E}\Delta_{i}\Delta_{j}-\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\right|\leq
K_{2}(x)\eta^{2},\quad\forall i,j=1,\ldots,d$; 3.
$\left|\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}-\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\right|\leq
K_{3}(x)\eta^{2},\quad\forall s\geq 3,\quad\forall i_{j}\in\\{1,\ldots,d\\}$;
4. $\mathbb{E}\prod_{j=1}^{3}\left|\bar{\Delta}_{i_{j}}\right|\leq
K_{4}(x)\eta^{2},\quad\forall i_{j}\in\\{1,\ldots,d\\}$. Then, there exists a
constant $C$ so that for all $k=0,1,\ldots,N$ we have
$\left|\mathbb{E}g\left(X_{k\eta}\right)-\mathbb{E}g\left(x_{k}\right)\right|\leq
C\eta.$

Before starting, we remind the reader that the order 1 weak approximation SDE
of SGD (see Li et al. (2017, 2019)) is given by

$dX_{t}=-\nabla
f(X_{t})dt+\sqrt{\eta}\left(\Sigma^{\text{SGD}}(X_{t})\right)^{\frac{1}{2}}dW_{t}$
(18)

where $\Sigma^{\text{SGD}}(x)$ is the SGD covariance matrix defined as

$\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)^{T}\right].$ (19)

###  A.1 Formal Derivation - USAM

The next result is inspired by Theorem 1 of (Li et al., 2017) and is derived
under some regularity assumption on the function $f$.

######  Assumption A.3.

Assume that the following conditions on $f,f_{i}$ and their gradients are
satisfied: • $\nabla f,\nabla f_{i}$ satisfy a Lipschitz condition: there
exists $L>0$ such that $|\nabla f(x)-\nabla f(y)|+\sum_{i=1}^{n}\left|\nabla
f_{i}(x)-\nabla f_{i}(y)\right|\leq L|x-y|;$ • $f,f_{i}$ and its partial
derivatives up to order 7 belong to $G$; • $\nabla f,\nabla f_{i}$ satisfy a
growth condition: there exists $M>0$ such that $|\nabla
f(x)|+\sum_{i=1}^{n}\left|\nabla f_{i}(x)\right|\leq M(1+|x|);$

We will consider the stochastic process $X_{t}\in\mathbb{R}^{d}$ defined by

$dX_{t}=-\nabla\tilde{f}^{\text{USAM}}(X_{t})dt+\sqrt{\eta}\left(\Sigma^{\text{SGD}}(X_{t})+\rho\left(\tilde{\Sigma}(X_{t})+\tilde{\Sigma}(X_{t})^{\top}\right)\right)^{\frac{1}{2}}dW_{t}$
(20)

where

$\Sigma^{\text{SGD}}(x):=\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)^{T}\right].$

is the usual covariance of SGD, while

$\tilde{\Sigma}(x):=\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left(\mathbb{E}\left[\nabla^{2}f_{\gamma}(x)\nabla
f_{\gamma}(x)\right]-\nabla^{2}f_{\gamma}(x)\nabla
f_{\gamma}(x)\right)^{\top}\right]$ (21)

and

$\tilde{f}^{\text{USAM}}(x):=f(x)+\frac{\rho}{2}\mathbb{E}\left[\lVert\nabla
f_{\gamma}(x)\rVert^{2}_{2}\right].$

In the following, we will use the notation

$\Sigma^{\text{USAM}}(x):=\left(\Sigma^{\text{SGD}}(X_{t})+\rho\left(\tilde{\Sigma}(X_{t})+\tilde{\Sigma}(X_{t})^{\top}\right)\right)$
(22)

######  Theorem A.4 (Stochastic modified equations).

Let $0<\eta<1,T>0$ and set $N=\lfloor T/\eta\rfloor$. Let
$x_{k}\in\mathbb{R}^{d},0\leq k\leq N$ denote a sequence of USAM iterations
defined by Eq. (5). Additionally, let us take
$\rho=\mathcal{O}\left(\eta^{\frac{1}{2}}\right).$ (23) Consider the
stochastic process $X_{t}$ defined in Eq. (20) and fix some test function
$g\in G$ and suppose that $g$ and its partial derivatives up to order 6 belong
to $G$. Then, under Assumption A.3, there exists a constant $C>0$ independent
of $\eta$ such that for all $k=0,1,\ldots,N$, we have
$\left|\mathbb{E}g\left(X_{k\eta}\right)-\mathbb{E}g\left(x_{k}\right)\right|\leq
C\eta^{1}.$ That is, the SDE (20) is an order $1$ weak approximation of the
SAM iterations (5).

######  Lemma A.5.

Under the assumptions of Theorem A.4, let $0<\eta<1$ and consider $x_{k},k\geq
0$ satisfying the USAM iterations (5) $x_{k+1}=x_{k}-\eta\nabla
f_{\gamma_{k}}\left(x_{k}+\rho\nabla f_{\gamma_{k}}(x_{k})\right)$ with
$x_{0}=x\in\mathbb{R}^{d}$. Additionally, we define
$\partial_{e_{i}}\tilde{f}^{\text{USAM}}(x):=\partial_{e_{i}}f(x)+\rho\mathbb{E}\left[\sum_{j}\partial_{e_{i}+e_{j}}^{2}f_{\gamma}(x)\partial_{e_{j}}f_{\gamma}(x)\right]$.
From the definition the one-step difference $\bar{\Delta}=x_{1}-x$, then we
have 1.
$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}\tilde{f}^{\text{USAM}}(x)\eta+\mathcal{O}(\eta\rho^{2})\quad\forall
i=1,\ldots,d$. 2.
$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=\partial_{e_{i}}\tilde{f}^{\text{USAM}}(x)\partial_{e_{j}}\tilde{f}^{\text{USAM}}(x)\eta^{2}+\Sigma^{\text{USAM}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)\quad\forall
i,j=1,\ldots,d$. 3.
$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{3}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}.$ and all the functions above are
evaluated at $x$.

Together with many other proofs in this Section, the following one relies on a
Taylor expansion. The truncated terms are multiplied by a mixture of terms in
$\eta$ and $\rho$. Therefore, a careful balancing of the relative size of
these two quantities is needed as reflected in Equation (23).

###### Proof of Lemma A.5.

Since the first step is to evaluate
$\mathbb{E}\Delta_{i}=\mathbb{E}\left[-\partial_{e_{i}}f_{\gamma}(x+\rho\nabla
f_{\gamma}(x))\eta\right]$, we start by analyzing
$\partial_{e_{i}}f_{\gamma}(x+\rho\nabla f_{\gamma}(x))$, that is the partial
derivative in the direction
$e_{i}:=(0,\cdots,0,\underset{i-th}{1},0,\cdots,0)$. Then, we have that

$\partial_{e_{i}}f_{\gamma}(x+\rho\nabla
f_{\gamma}(x))=\partial_{e_{i}}f_{\gamma}(x)+\sum_{|\alpha|=1}\partial_{e_{i}+\alpha}^{2}f_{\gamma}(x)\rho\partial_{\alpha}f_{\gamma}(x)+\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,1}(\rho\nabla
f_{\gamma}(x)),$ (24)

where the residual is defined in Eq. (4) of (Folland, 2005). Therefore, for
some constant $c\in(0,1)$, it holds that

$\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,1}(\rho\nabla
f_{\gamma}(x))=\sum_{|\alpha|=2}\frac{\partial_{e_{i}+\alpha}^{3}f_{\gamma}(x+c\rho\nabla
f_{\gamma}(x))\rho^{2}\left(\nabla f_{\gamma}(x)\right)^{\alpha}}{\alpha!}.$
(25)

Therefore, we can rewrite it as

$\displaystyle\partial_{e_{i}}f_{\gamma}(x+\rho\nabla
f_{\gamma}(x))=\partial_{e_{i}}f_{\gamma}(x)+\rho\sum_{j}\partial_{e_{i}+e_{j}}^{2}f_{\gamma}(x)\partial_{e_{j}}f_{\gamma}(x)+\rho^{2}\left[\sum_{|\alpha|=2}\frac{\partial_{e_{i}+\alpha}^{3}f_{\gamma}(x+c\rho\nabla
f_{\gamma}(x))\left(\nabla f_{\gamma}(x)\right)^{\alpha}}{\alpha!}\right]$
(26)

Now, we observe that

$K_{i}(x):=\left[\sum_{|\alpha|=2}\frac{\partial_{e_{i}+\alpha}^{3}f_{\gamma}(x+c\rho\nabla
f_{\gamma}(x))\left(\nabla f_{\gamma}(x)\right)^{\alpha}}{\alpha!}\right]$
(27)

is a finite sum of products of functions that by assumption are in $G$.
Therefore, $K_{i}(x)\in G$
and$\bar{K}_{i}(x)=\mathbb{E}\left[K_{i}(x)\right]\in G$. Based on these
definitions, we rewrite Eq. (26) as

$\partial_{e_{i}}f_{\gamma}(x+\rho\nabla
f_{\gamma}(x))=\partial_{e_{i}}f_{\gamma}(x)+\rho\sum_{j}\partial_{e_{i}+e_{j}}^{2}f_{\gamma}(x)\partial_{e_{j}}f_{\gamma}(x)+\rho^{2}K_{i}(x).$
(28)

which implies that

$\mathbb{E}\left[\partial_{e_{i}}f_{\gamma}(x+\rho\nabla
f_{\gamma}(x))\right]=\partial_{e_{i}}f(x)+\rho\mathbb{E}\left[\sum_{j}\partial_{e_{i}+e_{j}}^{2}f_{\gamma}(x)\partial_{e_{j}}f_{\gamma}(x)\right]+\rho^{2}\bar{K}_{i}(x).$
(29)

Let us now remember that

$\partial_{e_{i}}\tilde{f}^{\text{USAM}}(x)=\partial_{e_{i}}\left(f(x)+\frac{\rho}{2}\mathbb{E}\left[\lVert\nabla
f_{\gamma}(x)\rVert_{2}^{2}\right]\right)=\partial_{e_{i}}f(x)+\rho\mathbb{E}\left[\sum_{j}\partial_{e_{i}+e_{j}}^{2}f_{\gamma}(x)\partial_{e_{j}}f_{\gamma}(x)\right]$
(30)

Therefore, by using Eq. (29), Eq. (30), and the assumption (23) we have that
$\forall i=1,\ldots,d$

$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}\tilde{f}^{\text{USAM}}(x)\eta+\eta\rho^{2}\bar{K}_{i}(x)=-\partial_{e_{i}}\tilde{f}^{\text{USAM}}(x)\eta+\mathcal{O}\left(\eta^{2}\right).$
(31)

Additionally, we have that

$\displaystyle\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=$
$\displaystyle\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\mathbb{E}\bar{\Delta}_{i}\mathbb{E}\bar{\Delta}_{j}$
$\displaystyle\overset{\eqref{eq:USAM_Estim_Rewritten_3}}{=}\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\partial_{e_{i}}\tilde{f}^{\text{USAM}}\partial_{e_{j}}\tilde{f}^{\text{USAM}}\eta^{2}+\eta^{2}\rho^{2}(\partial_{e_{i}}\tilde{f}^{\text{USAM}}\bar{K}_{j}(x)+\partial_{e_{j}}\tilde{f}^{\text{USAM}}\bar{K}_{i}(x))+\eta^{2}\rho^{4}\bar{K}_{i}(x)\bar{K}_{j}(x)$
$\displaystyle=\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\partial_{e_{i}}\tilde{f}^{\text{USAM}}\partial_{e_{j}}\tilde{f}^{\text{USAM}}\eta^{2}+\mathcal{O}\left(\eta^{2}\rho^{2}\right)+\mathcal{O}\left(\eta^{2}\rho^{4}\right)$
$\displaystyle=\partial_{e_{i}}\tilde{f}^{\text{USAM}}\partial_{e_{j}}\tilde{f}^{\text{USAM}}\eta^{2}+\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\mathcal{O}\left(\eta^{2}\rho^{2}\right)+\mathcal{O}\left(\eta^{2}\rho^{4}\right)\quad\forall
i,j=1,\ldots,d$ (32)

Let us now recall the expression (21) of $\tilde{\Sigma}$ and the expression
(22) of $\Sigma^{\text{USAM}}$. Then, we automatically have that

$\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})=\eta^{2}\left(\Sigma^{\text{SGD}}_{i,j}(x)+\rho\left[\tilde{\Sigma}_{i,j}(x)+\tilde{\Sigma}_{i,j}(x)^{\top}\right]+\mathcal{O}(\rho^{2})\right)=\eta^{2}\Sigma^{\text{USAM}}_{i,j}(x)+\mathcal{O}(\eta^{2}\rho^{2})$
(33)

Therefore, remembering Eq. (32) and Eq. (23) we have

$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=\partial_{e_{i}}\tilde{f}^{\text{USAM}}\partial_{e_{j}}\tilde{f}^{\text{USAM}}\eta^{2}+\Sigma^{\text{USAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right),\quad\forall
i,j=1,\ldots,d$ (34)

Finally, with analogous considerations, it is obvious that under our
assumptions

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{s}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}$

which in particular implies that

$\mathbb{E}\prod_{j=1}^{3}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{3}\right),\quad
i_{j}\in\\{1,\ldots,d\\}.$

∎

##### Additional Insights from Lemma A.5.

Let us notice that $\nabla f_{\gamma}(x)$ is dominated by a factor $M(1+|x|)$,
if all $\partial_{e_{i}+\alpha}^{3}f_{\gamma}(x)$ are limited by a common
constant $L$, for some positive constant $C$ we have that

$\displaystyle\left|K_{i}(x)\right|=$
$\displaystyle\rho^{2}\left|\sum_{|\alpha|=2}\frac{\partial_{e_{i}+\alpha}^{3}f_{\gamma}(x+c\rho\nabla
f_{\gamma}(x))\left(\nabla f_{\gamma}(x)\right)^{\alpha}}{\alpha!}\right|$
(35) $\displaystyle\leq\rho^{2}CL\lVert\nabla f_{\gamma}(x)\nabla
f_{\gamma}(x)^{\top}\rVert_{F}^{2}\leq\rho^{2}CLd^{2}M^{2}(1+\lvert
x\rvert)^{2}$ (36)

Therefore, $K_{i}(x)$ does not only lay in $G$, but has at most quadratic
growth.

###### Proof of Theorem A.4.

To prove this result, all we need to do is check the conditions in Theorem
A.2. As we apply Lemma A.1, we make the following choices:

  * •

$b(x)=-\nabla\tilde{f}^{\text{USAM}}\left(x\right)$;

  * •

$\sigma(x)=\Sigma^{\text{USAM}}(x)^{\frac{1}{2}}$.

First of all, we notice that $\forall i=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\overset{\text{1. Lemma
\ref{lemma:USAM_SDE}}}{=}-\partial_{e_{i}}\tilde{f}^{\text{USAM}}(x)\eta+\mathcal{O}(\eta^{2})$;

  * •

$\mathbb{E}\Delta_{i}\overset{\text{1. Lemma
\ref{lemma:li1}}}{=}-\partial_{e_{i}}\tilde{f}^{\text{USAM}}(x)\eta+\mathcal{O}\left(\eta^{2}\right)$.

Therefore, we have that for some $K_{1}(x)\in G$

$\left|\mathbb{E}\Delta_{i}-\mathbb{E}\bar{\Delta}_{i}\right|\leq
K_{1}(x)\eta^{2},\quad\forall i=1,\ldots,d.$ (37)

Additionally,we notice that $\forall i,j=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\overset{\text{2. Lemma
\ref{lemma:USAM_SDE}}}{=}\partial_{e_{i}}\tilde{f}^{\text{USAM}}\partial_{e_{j}}\tilde{f}^{\text{USAM}}\eta^{2}+\Sigma^{\text{USAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\Delta_{i}\Delta_{j}\overset{\text{2. Lemma
\ref{lemma:li1}}}{=}\partial_{e_{i}}\tilde{f}^{\text{USAM}}\partial_{e_{j}}\tilde{f}^{\text{USAM}}\eta^{2}+\Sigma^{\text{USAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{2}(x)\in G$

$\left|\mathbb{E}\Delta_{i}\Delta_{j}-\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\right|\leq
K_{2}(x)\eta^{2},\quad\forall i,j=1,\ldots,d$ (38)

Additionally, we notice that $\forall s\geq 3,\forall
i_{j}\in\\{1,\ldots,d\\}$, it holds that

  * •

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:USAM_SDE}}}{=}\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:li1}}}{=}\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{3}(x)\in G$

$\left|\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}-\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\right|\leq
K_{3}(x)\eta^{2}.$ (39)

Additionally, for some $K_{4}(x)\in G$, $\forall i_{j}\in\\{1,\ldots,d\\}$

$\mathbb{E}\prod_{j=1}^{3}\left|\bar{\Delta}_{\left(i_{j}\right)}\right|\overset{\text{3.
Lemma \ref{lemma:USAM_SDE}}}{\leq}K_{4}(x)\eta^{2}.$ (40)

Finally, Eq. (37), Eq. (38), Eq. (39), and Eq. (40) allow us to conclude the
proof.

∎

######  Corollary A.6.

Let us take the same assumptions of Theorem A.4. Additionally, let us assume
that the dynamics is near the minimizer. In this case, the noise structure is
such that the stochastic gradient can be written as $\nabla
f_{\gamma}(x)=\nabla f(x)+Z$ such that $Z$ is the noise that does not depend
on $x$. Therefore, the SDE (20) becomes
$dX_{t}=-\nabla\tilde{f}^{\text{USAM}}(X_{t})dt+\left(I_{d}+\rho\nabla^{2}f(X_{t})\right)\left(\eta\Sigma^{\text{SGD}}(X_{t})\right)^{\frac{1}{2}}dW_{t}$
(41) where $\Sigma^{\text{SGD}}(x):=\mathbb{E}\left[\left(\nabla
f\left(x\right)-\nabla f_{\gamma}\left(x\right)\right)\left(\nabla
f\left(x\right)-\nabla f_{\gamma}\left(x\right)\right)^{T}\right]$ is the
usual covariance of SGD, and
$\tilde{f}^{\text{USAM}}(x)=f(x)+\frac{\rho}{2}\lVert\nabla
f(x)\rVert^{2}_{2}.$

###### Proof of Corollary A.6.

Based on our assumption on the noise structure, we can rewrite Eq. (21) of the
matrix $\tilde{\Sigma}$ as

$\displaystyle\tilde{\Sigma}(x)$ $\displaystyle=\mathbb{E}\left[\left(\nabla
f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left(\mathbb{E}\left[\nabla^{2}f_{\gamma}(x)\nabla
f_{\gamma}(x)\right]-\nabla^{2}f_{\gamma}(x)\nabla
f_{\gamma}(x)\right)^{\top}\right]$ (42)
$\displaystyle=\nabla^{2}f(x)\mathbb{E}\left[\left(\nabla
f\left(x\right)-\nabla f_{\gamma}\left(x\right)\right)\left(\nabla f(x)-\nabla
f_{\gamma}(x)\right)^{\top}\right]$ (43)

Therefore, the Eq. (22) of the covariance $\Sigma^{\text{USAM}}$ becomes

$\Sigma^{\text{USAM}}(x)=\left(I_{d}+2\rho\nabla^{2}f(x)\right)\Sigma^{\text{SGD}}(X_{t})$
(44)

which implies that

$(\Sigma^{\text{USAM}}(x))^{\frac{1}{2}}\approx\left(I_{d}+\rho\nabla^{2}f(x)\right)(\Sigma^{\text{SGD}}(X_{t}))^{\frac{1}{2}}.$
(45)

Finally, we have that

$\displaystyle\tilde{f}^{\text{USAM}}(x)$
$\displaystyle:=f(x)+\frac{\rho}{2}\mathbb{E}\left[\lVert\nabla
f_{\gamma}(x)\rVert^{2}_{2}\right]=f(x)+\frac{\rho}{2}\mathbb{E}\left[\lVert\nabla
f(x)\rVert^{2}_{2}+Z^{2}+2Z\nabla f_{\gamma}(x)\right]$ (46)
$\displaystyle=f(x)+\frac{\rho}{2}\lVert\nabla
f(x)\rVert^{2}_{2}+\frac{\rho}{2}\mathbb{E}\left[Z^{2}\right]$ (47)

Since the component $\mathbb{E}\left[Z^{2}\right]$ is independent on $x$, we
ignore it and conclude that

$\tilde{f}^{\text{USAM}}(x)=f(x)+\frac{\rho}{2}\lVert\nabla
f(x)\rVert^{2}_{2}.$

∎

####  A.1.1 USAM is SGD if $\rho=\mathcal{O}(\eta)$

The following result is inspired by Theorem 1 of (Li et al., 2017).

######  Theorem A.7 (Stochastic modified equations).

Let $0<\eta<1,T>0$ and set $N=\lfloor T/\eta\rfloor$. Let
$x_{k}\in\mathbb{R}^{d},0\leq k\leq N$ denote a sequence of USAM iterations
defined by Eq. (5). Additionally, let us take
$\rho=\mathcal{O}\left(\eta^{1}\right).$ (48) Define $X_{t}\in\mathbb{R}^{d}$
as the stochastic process satisfying the SDE $dX_{t}=-\nabla
f\left(X_{t}\right)dt+\left(\eta\Sigma^{\text{SGD}}\left(X_{t}\right)\right)^{1/2}dW_{t}$
(49) Such that $X_{0}=x_{0}$ and
$\Sigma^{\text{SGD}}(x):=\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)^{T}\right]$ Fix some test function $g\in G$
and suppose that $g$ and its partial derivatives up to order 6 belong to $G$.
Then, under Assumption A.3, there exists a constant $C>0$ independent of
$\eta$ such that for all $k=0,1,\ldots,N$, we have
$\left|\mathbb{E}g\left(X_{k\eta}\right)-\mathbb{E}g\left(x_{k}\right)\right|\leq
C\eta^{1}.$ That is, the SDE (49) is an order $1$ weak approximation of the
USAM iterations (5).

######  Lemma A.8.

Under the assumptions of Theorem A.7, let $0<\eta<1$. Consider $x_{k},k\geq 0$
satisfying the USAM iterations $x_{k+1}=x_{k}-\eta\nabla
f_{\gamma_{k}}\left(x_{k}+\rho\nabla f_{\gamma_{k}}(x_{k})\right)$ with
$x_{0}=x\in\mathbb{R}^{d}$. From the definition the one-step difference
$\bar{\Delta}=x_{1}-x$, then we have 1.
$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}f(x)\eta+\mathcal{O}(\eta^{2})\quad\forall
i=1,\ldots,d$. 2.
$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)\quad\forall
i,j=1,\ldots,d$. 3.
$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{3}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}.$ All functions above are evaluated at
$x$.

###### Proof of Lemma A.8.

First of all, we write that

$\partial_{e_{i}}f_{\gamma}(x+\rho\nabla
f_{\gamma}(x))=\partial_{e_{i}}f_{\gamma}(x)+\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,0}(\rho\nabla
f_{\gamma}(x)),$ (50)

where the residual is defined in Eq. (4) of (Folland, 2005). Therefore, for
some constant $c\in(0,1)$, it holds that

$\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,0}(\rho\nabla
f_{\gamma}(x))=\sum_{|\alpha|=1}\frac{\partial_{e_{i}+\alpha}^{2}f_{\gamma}(x+c\rho\nabla
f_{\gamma}(x))\rho^{1}\left(\nabla f_{\gamma}(x)\right)^{\alpha}}{\alpha!}.$
(51)

Let us now observe that
$\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,0}(\rho\nabla f_{\gamma}(x))$
is a finite sum of products of functions in $G$ and that, therefore, it lies
in $G$. Additionally, given its expression Eq. (51), we can factor out a
common $\rho$ and have that $K(x)=\rho K_{1}(x)$ for some function
$K_{1}(x)\in G$. Therefore, we rewrite Eq. (50) as

$\partial_{e_{i}}f_{\gamma}(x+\rho\nabla
f_{\gamma}(x))=\partial_{e_{i}}f_{\gamma}(x)+\rho K_{1}(x).$ (52)

First of all, we notice that if we define
$\bar{K}_{1}(x)=\mathbb{E}\left[K_{1}(x)\right]$, also $\bar{K}_{1}(x)\in G$.
Therefore, it holds that

$\mathbb{E}\left[\partial_{e_{i}}f_{\gamma}(x+\rho\nabla
f_{\gamma}(x))\right]\overset{\eqref{eq:USAM_SGD_Estim_Rewritten}}{=}\partial_{e_{i}}f(x)+\rho\bar{K}_{1}(x)$
(53)

Therefore, using assumption (48), $\forall i=1,\ldots,d$, we have that

$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}f(x)\eta+\eta\rho\bar{K}_{i}(x)=-\partial_{e_{i}}f(x)\eta+\mathcal{O}\left(\eta^{2}\right)$
(54)

Additionally, by keeping in mind the definition of the covariance matrix
$\Sigma$, We immediately have

$\displaystyle\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\overset{\eqref{eq:USAM_SGD_Estim_Rewritten}}{=}$
$\displaystyle
Cov(\bar{\Delta}_{i},\bar{\Delta}_{j})+\mathbb{E}\bar{\Delta}_{i}\mathbb{E}\bar{\Delta}_{j}$
$\displaystyle=\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\eta^{2}\rho(\partial_{e_{i}}f\bar{K}_{j}(x)+\partial_{e_{j}}f\bar{K}_{i}(x))+\eta^{2}\rho^{2}\bar{K}_{i}(x)\bar{K}_{j}(x)$
$\displaystyle=\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\mathcal{O}\left(\eta^{2}\rho\right)+\mathcal{O}\left(\eta^{2}\rho^{2}\right)$
$\displaystyle=\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)\quad\forall
i,j=1,\ldots,d$ (55)

Finally, with analogous considerations, it is obvious that under our
assumptions

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{3}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}.$

∎

###### Proof of Theorem A.7.

To prove this result, all we need to do is check the conditions in Theorem
A.2. As we apply Lemma A.1, we make the following choices:

  * •

$b(x)=-\nabla f\left(x\right)$,

  * •

$\sigma(x)=\Sigma^{\text{SGD}}(X_{t})^{\frac{1}{2}}$;

First of all, we notice that $\forall i=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\overset{\text{1. Lemma
\ref{lemma:USAM_SGD}}}{=}-\partial_{e_{i}}f(x)\eta+\mathcal{O}\left(\eta^{2}\right)$;

  * •

$\mathbb{E}\Delta_{i}\overset{\text{1. Lemma
\ref{lemma:li1}}}{=}-\partial_{e_{i}}f(x)\eta+\mathcal{O}\left(\eta^{2}\right)$.

Therefore, we have that for some $K_{1}(x)\in G$

$\left|\mathbb{E}\Delta_{i}-\mathbb{E}\bar{\Delta}_{i}\right|\leq
K_{1}(x)\eta^{2},\quad\forall i=1,\ldots,d.$ (56)

Additionally,we notice that $\forall i,j=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\overset{\text{2. Lemma
\ref{lemma:USAM_SGD}}}{=}\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\Delta_{i}\Delta_{j}\overset{\text{2. Lemma
\ref{lemma:li1}}}{=}\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{2}(x)\in G$

$\left|\mathbb{E}\Delta_{i}\Delta_{j}-\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\right|\leq
K_{2}(x)\eta^{2},\quad\forall i,j=1,\ldots,d$ (57)

Additionally, we notice that $\forall s\geq 3,\forall
i_{j}\in\\{1,\ldots,d\\}$, it holds that

  * •

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:USAM_SGD}}}{=}\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:li1}}}{=}\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{3}(x)\in G$

$\left|\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}-\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\right|\leq
K_{3}(x)\eta^{2}.$ (58)

Additionally, for some $K_{4}(x)\in G$, $\forall i_{j}\in\\{1,\ldots,d\\}$

$\mathbb{E}\prod_{j=1}^{3}\left|\bar{\Delta}_{\left(i_{j}\right)}\right|\overset{\text{3.
Lemma \ref{lemma:USAM_SGD}}}{\leq}K_{4}(x)\eta^{2}.$ (59)

Finally, Eq. (56), Eq. (57), Eq. (58), and Eq. (59) allow us to conclude the
proof.

∎

###  A.2 Formal Derivation - DNSAM

We now derive an SDE model for the DNSAM iteration given in (6) which we prove
to be a 1-order weak approximation of such a discrete iteration. The following
result is inspired by Theorem 1 of (Li et al., 2017). We will consider the
stochastic process $X_{t}\in\mathbb{R}^{d}$ defined as the solution of the SDE

$dX_{t}=-\nabla\tilde{f}^{\text{DNSAM}}(X_{t})dt+\left(I_{d}+\rho\frac{\nabla^{2}f(X_{t})}{\lVert\nabla
f(X_{t})\rVert_{2}}\right)\left(\eta\Sigma^{\text{SGD}}(X_{t})\right)^{\frac{1}{2}}dW_{t}$
(60)

where the regularized loss is

$\tilde{f}^{\text{DNSAM}}(x)=f(x)+\rho\lVert\nabla f(x)\rVert_{2},$

the covariance matrix is

$\Sigma^{DNSAM}(x):=\Sigma^{SGD}(x)\left(I_{d}+2\rho\frac{\nabla^{2}f(x)}{\lVert\nabla
f(x)\rVert}\right)$ (61)

and

$\Sigma^{\text{SGD}}(x):=\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)^{T}\right]$

is the usual covariance of SGD.

######  Theorem A.9 (Stochastic modified equations).

Let $0<\eta<1,T>0$ and set $N=\lfloor T/\eta\rfloor$. Let
$x_{k}\in\mathbb{R}^{d},0\leq k\leq N$ denote a sequence of DNSAM iterations
defined by Eq. (6). Additionally, let us assume that the noise structure is
such that the stochastic gradient can be written as $\nabla
f_{\gamma}(x)=\nabla f(x)+Z$ and
$\rho=\mathcal{O}\left(\eta^{\frac{1}{2}}\right).$ (62) Consider the
stochastic process $X_{t}$ defined in Eq. (60) and fix some test function
$g\in G$ and suppose that $g$ and its partial derivatives up to order 6 belong
to $G$. Then, under Assumption A.3, there exists a constant $C>0$ independent
of $\eta$ such that for all $k=0,1,\ldots,N$, we have
$\left|\mathbb{E}g\left(X_{k\eta}\right)-\mathbb{E}g\left(x_{k}\right)\right|\leq
C\eta^{1}.$ That is, the SDE (60) is an order $1$ weak approximation of the
DNSAM iterations (6).

######  Lemma A.10.

Under the assumptions of Theorem A.9, let $0<\eta<1$ and consider $x_{k},k\geq
0$ satisfying the DNSAM iterations (6) $x_{k+1}=x_{k}-\eta\nabla
f_{\gamma_{k}}\left(x_{k}+\rho\frac{\nabla f_{\gamma_{k}}(x_{k})}{\lVert\nabla
f(x_{k})\rVert}\right)$ with $x_{0}=x\in\mathbb{R}^{d}$. Additionally, we
define
$\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}(x):=\partial_{e_{i}}f(x)+\rho\frac{\sum_{j}\partial_{e_{i}+e_{j}}^{2}f(x)\partial_{e_{j}}f(x)}{\lVert\nabla
f(x)\rVert}$. From the definition the one-step difference
$\bar{\Delta}=x_{1}-x$, and we indicate with $\bar{\Delta}_{i}$ the $i$-th
component of such difference. Then, we have 1.
$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}(x)\eta+\mathcal{O}(\eta\rho^{2})\quad\forall
i=1,\ldots,d$; 2.
$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}(x)\partial_{e_{j}}\tilde{f}^{\text{DNSAM}}(x)\eta^{2}+\Sigma^{\text{DNSAM}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)\quad\forall
i,j=1,\ldots,d$; 3.
$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta_{i_{j}}}=\mathcal{O}\left(\eta^{3}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}.$ and all the functions above are
evaluated at $x$.

###### Proof of Lemma A.10.

Since the first step is to evaluate
$\mathbb{E}\Delta_{i}=-\mathbb{E}\left[\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla
f(x)\rVert}\nabla f_{\gamma}(x)\right)\eta\right]$, we start by analyzing
$\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla f(x)\rVert}\nabla
f_{\gamma}(x)\right)$, that is the partial derivative in the direction
$e_{i}:=(0,\cdots,0,\underset{i-th}{1},0,\cdots,0)$. Under the noise
assumption $\nabla f_{\gamma}(x)=\nabla f(x)+Z$, we have that
$\nabla^{2}f_{\gamma}(x)=\nabla^{2}f(x)$. Then, we have that

$\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla f(x)\rVert}\nabla
f_{\gamma}(x)\right)=\partial_{e_{i}}f_{\gamma}(x)+\sum_{|\alpha|=1}\partial_{e_{i}+\alpha}^{2}f(x)\rho\frac{\partial_{\alpha}f_{\gamma}(x)}{\lVert\nabla
f(x)\rVert}+\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,1}\left(\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f(x)\rVert}\right),$ (63)

where the residual is defined in Eq. (4) of (Folland, 2005). Therefore, for
some constant $c\in(0,1)$, it holds that

$\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,1}\left(\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla
f(x)\rVert}\right)=\sum_{|\alpha|=2}\frac{\partial_{e_{i}+\alpha}^{3}f_{\gamma}\left(x+c\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f(x)\rVert}\right)\rho^{2}\left(\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f(x)\rVert}\right)^{\alpha}}{\alpha!}.$ (64)

Combining the last two equations, we obtain

$\displaystyle\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla
f(x)\rVert}\nabla f_{\gamma}(x)\right)$
$\displaystyle=\partial_{e_{i}}f_{\gamma}(x)+\frac{\rho}{\lVert\nabla
f(x)\rVert}\sum_{|\alpha|=1}\partial_{e_{i}+\alpha}^{2}f(x)\partial_{\alpha}f_{\gamma}(x)$
(65)
$\displaystyle+\rho^{2}\sum_{|\alpha|=2}\frac{\partial_{e_{i}+\alpha}^{3}f_{\gamma}\left(x+c\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f(x)\rVert}\right)\left(\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f(x)\rVert}\right)^{\alpha}}{\alpha!}.$ (66)

Now, we observe that

$K_{i}(x):=\left[\sum_{|\alpha|=2}\frac{\partial_{e_{i}+\alpha}^{3}f_{\gamma}\left(x+c\rho\frac{\nabla
f(x)}{\lVert\nabla f(x)\rVert}\right)\left(\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f(x)\rVert}\right)^{\alpha}}{\alpha!}\right]$ (67)

is a finite sum of products of functions that by assumption are in $G$.
Therefore, $K_{i}(x)\in G$ and
$\bar{K}_{i}(x)=\mathbb{E}\left[K_{i}(x)\right]\in G$. Based on these
definitions, we rewrite Eq. (65) as

$\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla f(x)\rVert}\nabla
f_{\gamma}(x)\right)=\partial_{e_{i}}f_{\gamma}(x)+\frac{\rho}{\lVert\nabla
f(x)\rVert}\sum_{|\alpha|=1}\partial_{e_{i}+\alpha}^{2}f(x)\partial_{\alpha}f_{\gamma}(x)+\rho^{2}K_{i}(x).$
(68)

which implies that

$\mathbb{E}\left[\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla
f(x)\rVert}\nabla
f_{\gamma}(x)\right)\right]=\partial_{e_{i}}f(x)+\rho\frac{\sum_{j}\partial_{e_{i}+e_{j}}^{2}f(x)\partial_{e_{j}}f(x)}{\lVert\nabla
f(x)\rVert}+\rho^{2}\bar{K}_{i}(x),$ (69)

where we used the unbiasedness property of the stochastic gradients:
${\mathbb{E}}\nabla f_{\gamma}(x)=\nabla f(x)$.

Let us now remember that by definition

$\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}(x)=\partial_{e_{i}}f(x)+\rho\frac{\sum_{j}\partial_{e_{i}+e_{j}}^{2}f(x)\partial_{e_{j}}f(x)}{\lVert\nabla
f(x)\rVert}.$ (70)

Therefore, by using Eq. (69), Eq. (70), and the assumption (62) we have that
$\forall i=1,\ldots,d$,

$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}(x)\eta+\eta\rho^{2}\bar{K}_{i}(x)=-\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}(x)\eta+\mathcal{O}\left(\eta^{2}\right).$
(71)

We now observe that the covariance matrix of the difference between the drift
$\nabla f(x)+\rho\frac{\nabla^{2}f(x)\nabla f(x)}{\lVert\nabla f(x)\rVert}$ of
the SDE (60) and the gradient $\nabla
f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla f(x)\rVert}\nabla
f_{\gamma}(x)\right)=\nabla f_{\gamma}(x)+\rho\frac{\nabla^{2}f(x)\nabla
f_{\gamma}(x)}{\lVert\nabla f(x)\rVert}+\rho^{2}K(x)$ in the discrete
algorithm (6) is

$\displaystyle\bar{\Sigma}:=\mathbb{E}$ $\displaystyle\left[\left(\nabla
f(x)+\rho\frac{\nabla^{2}f(x)\nabla f(x)}{\lVert\nabla f(x)\rVert}-\nabla
f_{\gamma}(x)-\rho\frac{\nabla^{2}f(x)\nabla f_{\gamma}(x)}{\lVert\nabla
f(x)\rVert}-\rho^{2}K(x)\right)\right.$ (72) $\displaystyle\left.\left(\nabla
f(x)+\rho\frac{\nabla^{2}f(x)\nabla f(x)}{\lVert\nabla f(x)\rVert}-\nabla
f_{\gamma}(x)-\rho\frac{\nabla^{2}f(x)\nabla f_{\gamma}(x)}{\lVert\nabla
f(x)\rVert}-\rho^{2}K(x)\right)^{\top}\right]$ (73)
$\displaystyle=\Sigma^{SGD}\left(I_{d}+2\rho\frac{\nabla^{2}f(x)}{\lVert\nabla
f(x)\rVert}\right)+\mathcal{O}(\rho^{2})=\Sigma^{DNSAM}(x)+\mathcal{O}(\rho^{2}).$
(74)

Therefore, we have that

$\displaystyle\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=$
$\displaystyle\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\mathbb{E}\bar{\Delta}_{i}\mathbb{E}\bar{\Delta}_{j}$
$\displaystyle\overset{\eqref{eq:DNSAM_Estim_Rewritten_3}}{=}\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}\partial_{e_{j}}\tilde{f}^{\text{DNSAM}}\eta^{2}+\eta^{2}\rho^{2}(\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}\bar{K}_{j}(x)+\partial_{e_{j}}\tilde{f}^{\text{DNSAM}}\bar{K}_{i}(x))+\eta^{2}\rho^{4}\bar{K}_{i}(x)\bar{K}_{j}(x)$
$\displaystyle=\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}\partial_{e_{j}}\tilde{f}^{\text{DNSAM}}\eta^{2}+\mathcal{O}\left(\eta^{2}\rho^{2}\right)+\mathcal{O}\left(\eta^{2}\rho^{4}\right)$
$\displaystyle=\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}\partial_{e_{j}}\tilde{f}^{\text{DNSAM}}\eta^{2}+\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\mathcal{O}\left(\eta^{2}\rho^{2}\right)+\mathcal{O}\left(\eta^{2}\rho^{4}\right)\quad\forall
i,j=1,\ldots,d.$ (75)

By the definitions of $\Sigma^{\text{DNSAM}}$ and of $\bar{\Sigma}(x)$, we
have

$\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})=\eta^{2}\bar{\Sigma}_{i,j}(x)=\eta^{2}\Sigma^{\text{DNSAM}}_{i,j}(x)+\mathcal{O}(\eta^{2}\rho^{2}).$
(76)

Therefore, remembering Eq. (75) and Eq. (62) we have

$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}\partial_{e_{j}}\tilde{f}^{\text{DNSAM}}\eta^{2}+\Sigma^{\text{DNSAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right),\quad\forall
i,j=1,\ldots,d.$ (77)

Finally, with analogous considerations, it is obvious that under our
assumptions

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{s}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}$

which in particular implies that

$\mathbb{E}\prod_{j=1}^{3}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{3}\right),\quad
i_{j}\in\\{1,\ldots,d\\}.$

∎

###### Proof of Theorem A.9.

To prove this result, all we need to do is check the conditions in Theorem
A.2. As we apply Lemma A.1, we make the following choices:

  * •

$b(x)=-\nabla\tilde{f}^{\text{DNSAM}}\left(x\right)$;

  * •

$\sigma(x)=\Sigma^{\text{DNSAM}}(x)^{\frac{1}{2}}$.

First of all, we notice that $\forall i=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\overset{\text{1. Lemma
\ref{lemma:DNSAM_SDE}}}{=}-\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}(x)\eta+\mathcal{O}(\eta^{2})$;

  * •

$\mathbb{E}\Delta_{i}\overset{\text{1. Lemma
\ref{lemma:li1}}}{=}-\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}(x)\eta+\mathcal{O}\left(\eta^{2}\right)$.

Therefore, we have that for some $K_{1}(x)\in G$

$\left|\mathbb{E}\Delta_{i}-\mathbb{E}\bar{\Delta}_{i}\right|\leq
K_{1}(x)\eta^{2},\quad\forall i=1,\ldots,d.$ (78)

Additionally,we notice that $\forall i,j=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\overset{\text{2. Lemma
\ref{lemma:DNSAM_SDE}}}{=}\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}\partial_{e_{j}}\tilde{f}^{\text{DNSAM}}\eta^{2}+\Sigma^{\text{DNSAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\Delta_{i}\Delta_{j}\overset{\text{2. Lemma
\ref{lemma:li1}}}{=}\partial_{e_{i}}\tilde{f}^{\text{DNSAM}}\partial_{e_{j}}\tilde{f}^{\text{DNSAM}}\eta^{2}+\Sigma^{\text{DNSAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{2}(x)\in G$

$\left|\mathbb{E}\Delta_{i}\Delta_{j}-\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\right|\leq
K_{2}(x)\eta^{2},\quad\forall i,j=1,\ldots,d.$ (79)

Additionally, we notice that $\forall s\geq 3,\forall
i_{j}\in\\{1,\ldots,d\\}$, it holds that

  * •

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:DNSAM_SDE}}}{=}\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:li1}}}{=}\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{3}(x)\in G$

$\left|\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}-\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\right|\leq
K_{3}(x)\eta^{2}.$ (80)

Additionally, for some $K_{4}(x)\in G$, $\forall i_{j}\in\\{1,\ldots,d\\}$

$\mathbb{E}\prod_{j=1}^{3}\left|\bar{\Delta}_{\left(i_{j}\right)}\right|\overset{\text{3.
Lemma \ref{lemma:DNSAM_SDE}}}{\leq}K_{4}(x)\eta^{2}.$ (81)

Finally, Eq. (78), Eq. (79), Eq. (80), and Eq. (81) allow us to conclude the
proof.

∎

####  A.2.1 DNSAM is SGD if $\rho=\mathcal{O}(\eta)$

The following result is inspired by Theorem 1 of (Li et al., 2017). We will
consider the stochastic process $X_{t}\in\mathbb{R}^{d}$ defined as the
solution of the SDE

$dX_{t}=-\nabla
f\left(X_{t}\right)dt+\left(\eta\Sigma^{\text{SGD}}\left(X_{t}\right)\right)^{1/2}dW_{t}$
(82)

Such that $X_{0}=x_{0}$ and

$\Sigma^{\text{SGD}}(x):=\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)^{T}\right]$

######  Theorem A.11 (Stochastic modified equations).

Let $0<\eta<1,T>0$ and set $N=\lfloor T/\eta\rfloor$. Let
$x_{k}\in\mathbb{R}^{d},0\leq k\leq N$ denote a sequence of DNSAM iterations
defined by Eq. (6). Additionally, let us take
$\rho=\mathcal{O}\left(\eta^{1}\right).$ (83) Consider the stochastic process
$X_{t}$ defined in Eq. (82) and fix some test function $g\in G$ and suppose
that $g$ and its partial derivatives up to order 6 belong to $G$. Then, under
Assumption A.3, there exists a constant $C>0$ independent of $\eta$ such that
for all $k=0,1,\ldots,N$, we have
$\left|\mathbb{E}g\left(X_{k\eta}\right)-\mathbb{E}g\left(x_{k}\right)\right|\leq
C\eta^{1}.$ That is, the SDE (82) is an order $1$ weak approximation of the
SAM iterations (4).

###### Proof.

The proof is completely similar to that of Theorem A.7 presented before and of
Theorem A.16 presented later. ∎

###  A.3 Formal Derivation - SAM

The following result is inspired by Theorem 1 of (Li et al., 2017). We will
consider the stochastic process $X_{t}\in\mathbb{R}^{d}$ defined as the
solution of the SDE

$dX_{t}=-\nabla\tilde{f}^{\text{SAM}}(X_{t})dt+\sqrt{\eta}\left(\Sigma^{\text{SGD}}(X_{t})+\rho\left(\hat{\Sigma}(X_{t})+\hat{\Sigma}(X_{t})^{\top}\right)\right)^{\frac{1}{2}}dW_{t}$
(84)

where

$\Sigma^{\text{SGD}}(x):=\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)^{T}\right]$

is the usual covariance of SGD, while

$\hat{\Sigma}(x):=\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left({\mathbb{E}}\left[\frac{\nabla^{2}f_{\gamma}(x)\nabla
f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert_{2}}\right]-\frac{\nabla^{2}f_{\gamma}(x)\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert_{2}}\right)^{\top}\right]$
(85)

and

$\tilde{f}^{\text{SAM}}(x):=f(x)+\rho\mathbb{E}\left[\lVert\nabla
f_{\gamma}(x)\rVert_{2}\right].$

In the following, we will use the notation

$\Sigma^{\text{SAM}}(x):=\left(\Sigma^{\text{SGD}}(X_{t})+\rho\left(\hat{\Sigma}(X_{t})+\hat{\Sigma}(X_{t})^{\top}\right)\right).$
(86)

######  Theorem A.12 (Stochastic modified equations).

Let $0<\eta<1,T>0$ and set $N=\lfloor T/\eta\rfloor$. Let
$x_{k}\in\mathbb{R}^{d},0\leq k\leq N$ denote a sequence of SAM iterations
defined by Eq. (4). Additionally, let us take
$\rho=\mathcal{O}\left(\eta^{\frac{1}{2}}\right).$ (87) Consider the
stochastic process $X_{t}$ defined in Eq. (84) and fix some test function
$g\in G$ and suppose that $g$ and its partial derivatives up to order 6 belong
to $G$. Then, under Assumption A.3, there exists a constant $C>0$ independent
of $\eta$ such that for all $k=0,1,\ldots,N$, we have
$\left|\mathbb{E}g\left(X_{k\eta}\right)-\mathbb{E}g\left(x_{k}\right)\right|\leq
C\eta^{1}.$ That is, the SDE (84) is an order $1$ weak approximation of the
SAM iterations (4).

######  Remark A.13.

Denote by $b$ and $\sigma$ the (Borel measurable) drift and diffusion
coefficient in (84), respectively. Suppose that there exists a non-negative
function $F\in L^{d+1}(\left[0,\infty\right)\times\mathbb{R}^{d})$ such that

$\left\|b(t,x)\right\|\leq K+F(t,x)$ (88)

$dt\times dx-a.e.$ for some constant $K\geq 0$. Further, assume that
$\sigma\sigma^{\top}$ is strongly elliptic, that is there exists a
$\delta\in(0,1)$ such that for all $t\geq 0$, $x\in\mathbb{R}^{d}$

$\delta I_{d}\leq\sigma\sigma^{\top}(t,x)\leq\delta^{-1}I_{d},$ (89)

where $I_{d}\in\mathbb{R}^{d\times d}$ is the unit matrix. Then there exists a
(global) weak solution to (84). Moreover, if there is a (global) weak solution
to (84), $b\in L_{loc}^{2d+2}(\left[0,\infty\right)\times\mathbb{R}^{d})$,
$\sigma$ is locally Lipschitz continuous uniformly in time and if (89) holds,
then there exists a unique strong solution to (84). See (Gyöngy & Martínez,
2001).

######  Lemma A.14.

Under the assumptions of Theorem A.12, let $0<\eta<1$ and consider
$x_{k},k\geq 0$ satisfying the USAM iterations (4) $x_{k+1}=x_{k}-\eta\nabla
f_{\gamma_{k}}\left(x_{k}+\rho\frac{\nabla f_{\gamma_{k}}(x_{k})}{\lVert\nabla
f_{\gamma}(x_{k})\rVert}\right)$ with $x_{0}=x\in\mathbb{R}^{d}$.
Additionally, we define
$\partial_{e_{i}}\tilde{f}^{\text{SAM}}(x):=\partial_{e_{i}}f(x)+\rho\mathbb{E}\left[\frac{\sum_{j}\partial_{e_{i}+e_{j}}^{2}f_{\gamma}(x)\partial_{e_{j}}f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}\right]$. From the definition the one-step difference
$\bar{\Delta}=x_{1}-x$, then we have 1.
$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}\tilde{f}^{\text{SAM}}(x)\eta+\mathcal{O}(\eta\rho^{2})\quad\forall
i=1,\ldots,d$; 2.
$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=\partial_{e_{i}}\tilde{f}^{\text{SAM}}(x)\partial_{e_{j}}\tilde{f}^{\text{SAM}}(x)\eta^{2}+\Sigma^{\text{SAM}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)\quad\forall
i,j=1,\ldots,d$; 3.
$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{3}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}.$ and all the functions above are
evaluated at $x$.

###### Proof of Lemma A.14.

Since the first step is to evaluate
$\mathbb{E}\Delta_{i}=-\mathbb{E}\left[\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla
f_{\gamma}(x)\rVert}\nabla f_{\gamma}(x)\right)\eta\right]$, we start by
analyzing $\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla
f_{\gamma}(x)\rVert}\nabla f_{\gamma}(x)\right)$, that is the partial
derivative in the direction
$e_{i}:=(0,\cdots,0,\underset{i-th}{1},0,\cdots,0)$. Then, we have that

$\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla
f_{\gamma}(x)\rVert}\nabla
f_{\gamma}(x)\right)=\partial_{e_{i}}f_{\gamma}(x)+\sum_{|\alpha|=1}\partial_{e_{i}+\alpha}^{2}f_{\gamma}(x)\rho\frac{\partial_{\alpha}f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}+\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,1}\left(\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert}\right)$ (90)

Where the residual is defined in Eq. (4) of (Folland, 2005). Therefore, for
some constant $c\in(0,1)$, it holds that

$\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,1}\left(\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}\right)=\sum_{|\alpha|=2}\frac{\partial_{e_{i}+\alpha}^{3}f_{\gamma}\left(x+c\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}\right)\rho^{2}\left(\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert}\right)^{\alpha}}{\alpha!}.$
(91)

Therefore, we can rewrite it as

$\displaystyle\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla
f_{\gamma}(x)\rVert}\nabla f_{\gamma}(x)\right)$
$\displaystyle=\partial_{e_{i}}f_{\gamma}(x)+\frac{\rho}{\lVert\nabla
f_{\gamma}(x)\rVert}\sum_{|\alpha|=1}\partial_{e_{i}+\alpha}^{2}f_{\gamma}(x)\partial_{\alpha}f_{\gamma}(x)$
(92)
$\displaystyle+\rho^{2}\sum_{|\alpha|=2}\frac{\partial_{e_{i}+\alpha}^{3}f_{\gamma}\left(x+c\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert}\right)\left(\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert}\right)^{\alpha}}{\alpha!}.$
(93)

Now, we observe that

$K_{i}(x):=\left[\sum_{|\alpha|=2}\frac{\partial_{e_{i}+\alpha}^{3}f_{\gamma}\left(x+c\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert}\right)\left(\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}\right)^{\alpha}}{\alpha!}\right]$ (94)

is a finite sum of products of functions that by assumption are in $G$.
Therefore, $K_{i}(x)\in G$ and
$\bar{K}_{i}(x)=\mathbb{E}\left[K_{i}(x)\right]\in G$. Based on these
definitions, we rewrite Eq. (92) as

$\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla
f_{\gamma}(x)\rVert}\nabla
f_{\gamma}(x)\right)=\partial_{e_{i}}f_{\gamma}(x)+\frac{\rho}{\lVert\nabla
f_{\gamma}(x)\rVert}\sum_{|\alpha|=1}\partial_{e_{i}+\alpha}^{2}f_{\gamma}(x)\partial_{\alpha}f_{\gamma}(x)+\rho^{2}K_{i}(x).$
(95)

which implies that

$\mathbb{E}\left[\partial_{e_{i}}f_{\gamma}\left(x+\frac{\rho}{\lVert\nabla
f_{\gamma}(x)\rVert}\nabla
f_{\gamma}(x)\right)\right]=\partial_{e_{i}}f(x)+\rho\mathbb{E}\left[\frac{\sum_{j}\partial_{e_{i}+e_{j}}^{2}f_{\gamma}(x)\partial_{e_{j}}f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}\right]+\rho^{2}\bar{K}_{i}(x).$ (96)

Let us now remember that

$\partial_{e_{i}}\tilde{f}^{\text{SAM}}(x)=\partial_{e_{i}}\left(f(x)+\rho\mathbb{E}\left[\lVert\nabla
f_{\gamma}(x)\rVert_{2}\right]\right)=\partial_{e_{i}}f(x)+\rho\mathbb{E}\left[\frac{\sum_{j}\partial_{e_{i}+e_{j}}^{2}f_{\gamma}(x)\partial_{e_{j}}f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}\right]$ (97)

Therefore, by using Eq. (96), Eq. (97), and the assumption (87) we have that
$\forall i=1,\ldots,d$

$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}\tilde{f}^{\text{SAM}}(x)\eta+\eta\rho^{2}\bar{K}_{i}(x)=-\partial_{e_{i}}\tilde{f}^{\text{SAM}}(x)\eta+\mathcal{O}\left(\eta^{2}\right).$
(98)

Additionally, we have that

$\displaystyle\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=$
$\displaystyle\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\mathbb{E}\bar{\Delta}_{i}\mathbb{E}\bar{\Delta}_{j}$
$\displaystyle\overset{\eqref{eq:SAM_Estim_Rewritten_3}}{=}\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\partial_{e_{i}}\tilde{f}^{\text{SAM}}\partial_{e_{j}}\tilde{f}^{\text{SAM}}\eta^{2}+\eta^{2}\rho^{2}(\partial_{e_{i}}\tilde{f}^{\text{SAM}}\bar{K}_{j}(x)+\partial_{e_{j}}\tilde{f}^{\text{SAM}}\bar{K}_{i}(x))+\eta^{2}\rho^{4}\bar{K}_{i}(x)\bar{K}_{j}(x)$
$\displaystyle=\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\partial_{e_{i}}\tilde{f}^{\text{SAM}}\partial_{e_{j}}\tilde{f}^{\text{SAM}}\eta^{2}+\mathcal{O}\left(\eta^{2}\rho^{2}\right)+\mathcal{O}\left(\eta^{2}\rho^{4}\right)$
$\displaystyle=\partial_{e_{i}}\tilde{f}^{\text{SAM}}\partial_{e_{j}}\tilde{f}^{\text{SAM}}\eta^{2}+\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\mathcal{O}\left(\eta^{2}\rho^{2}\right)+\mathcal{O}\left(\eta^{2}\rho^{4}\right)\quad\forall
i,j=1,\ldots,d.$ (99)

Let us now recall the expression (85) of $\hat{\Sigma}$ and the expression
(86) of $\Sigma^{\text{SAM}}$. Then, we automatically have that

$\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})=\eta^{2}\left(\Sigma^{\text{SGD}}_{i,j}(x)+\rho\left[\hat{\Sigma}_{i,j}(x)+\hat{\Sigma}_{i,j}(x)^{\top}\right]+\mathcal{O}(\rho^{2})\right)=\eta^{2}\Sigma^{\text{SAM}}_{i,j}(x)+\mathcal{O}(\eta^{2}\rho^{2}).$
(100)

Therefore, remembering Eq. (99) and Eq. (87) we have

$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=\partial_{e_{i}}\tilde{f}^{\text{SAM}}\partial_{e_{j}}\tilde{f}^{\text{SAM}}\eta^{2}+\Sigma^{\text{SAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right),\quad\forall
i,j=1,\ldots,d.$ (101)

Finally, with analogous considerations, it is obvious that under our
assumptions

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{s}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}$

which in particular implies that

$\mathbb{E}\prod_{j=1}^{3}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{3}\right),\quad
i_{j}\in\\{1,\ldots,d\\}.$

∎

###### Proof of Theorem A.12.

To prove this result, all we need to do is check the conditions in Theorem
A.2. As we apply Lemma A.1, we make the following choices:

  * •

$b(x)=-\nabla\tilde{f}^{\text{SAM}}\left(x\right)$;

  * •

$\sigma(x)=\Sigma^{\text{SAM}}(x)^{\frac{1}{2}}$.

First of all, we notice that $\forall i=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\overset{\text{1. Lemma
\ref{lemma:SAM_SDE}}}{=}-\partial_{e_{i}}\tilde{f}^{\text{SAM}}(x)\eta+\mathcal{O}(\eta^{2})$;

  * •

$\mathbb{E}\Delta_{i}\overset{\text{1. Lemma
\ref{lemma:li1}}}{=}-\partial_{e_{i}}\tilde{f}^{\text{SAM}}(x)\eta+\mathcal{O}\left(\eta^{2}\right)$.

Therefore, we have that for some $K_{1}(x)\in G$

$\left|\mathbb{E}\Delta_{i}-\mathbb{E}\bar{\Delta}_{i}\right|\leq
K_{1}(x)\eta^{2},\quad\forall i=1,\ldots,d.$ (102)

Additionally,we notice that $\forall i,j=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\overset{\text{2. Lemma
\ref{lemma:SAM_SDE}}}{=}\partial_{e_{i}}\tilde{f}^{\text{SAM}}\partial_{e_{j}}\tilde{f}^{\text{SAM}}\eta^{2}+\Sigma^{\text{SAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\Delta_{i}\Delta_{j}\overset{\text{2. Lemma
\ref{lemma:li1}}}{=}\partial_{e_{i}}\tilde{f}^{\text{SAM}}\partial_{e_{j}}\tilde{f}^{\text{SAM}}\eta^{2}+\Sigma^{\text{SAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{2}(x)\in G$

$\left|\mathbb{E}\Delta_{i}\Delta_{j}-\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\right|\leq
K_{2}(x)\eta^{2},\quad\forall i,j=1,\ldots,d.$ (103)

Additionally, we notice that $\forall s\geq 3,\forall
i_{j}\in\\{1,\ldots,d\\}$, it holds that

  * •

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:SAM_SDE}}}{=}\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:li1}}}{=}\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{3}(x)\in G$

$\left|\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}-\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\right|\leq
K_{3}(x)\eta^{2}.$ (104)

Additionally, for some $K_{4}(x)\in G$, $\forall i_{j}\in\\{1,\ldots,d\\}$

$\mathbb{E}\prod_{j=1}^{3}\left|\bar{\Delta}_{\left(i_{j}\right)}\right|\overset{\text{3.
Lemma \ref{lemma:USAM_SDE}}}{\leq}K_{4}(x)\eta^{2}.$ (105)

Finally, Eq. (102), Eq. (103), Eq. (104), and Eq. (105) allow us to conclude
the proof.

∎

######  Corollary A.15.

Let us take the same assumptions of Theorem (A.12). Additionally, let us
assume that the dynamics is near the minimizer. In this case, the noise
structure is such that the stochastic gradient can be written as $\nabla
f_{\gamma}(x)=\nabla f(x)+Z$ such that $Z$ is the noise that does not depend
on $x$. In this case, the SDE (84) becomes
$dX_{t}=-\nabla\tilde{f}^{\text{SAM}}(X_{t})dt+\sqrt{\eta\left(\Sigma^{\text{SGD}}(X_{t})+\rho
H_{t}\left(\bar{\Sigma}(X_{t})+\bar{\Sigma}(X_{t})^{\top}\right)\right)}dW_{t}$
where $H_{t}:=\nabla^{2}f(X_{t})$ and $\bar{\Sigma}(x)$ is defined as
$\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left({\mathbb{E}}\left[\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert_{2}}\right]-\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert_{2}}\right)^{\top}\right],$
and $\tilde{f}^{\text{SAM}}(x):=f(x)+\rho\mathbb{E}\left[\lVert\nabla
f_{\gamma}(x)\rVert_{2}\right].$

###### Proof of Corollary A.15.

It follows immediately by substituting the expression for the perturbed
gradient.

∎

####  A.3.1 SAM is SGD if $\rho=\mathcal{O}(\eta)$

The following result is inspired by Theorem 1 of (Li et al., 2017). We will
consider the stochastic process $X_{t}\in\mathbb{R}^{d}$ defined as the
solution of the SDE

$dX_{t}=-\nabla
f\left(X_{t}\right)dt+\left(\eta\Sigma^{\text{SGD}}\left(X_{t}\right)\right)^{1/2}dW_{t}$
(106)

Such that $X_{0}=x_{0}$ and

$\Sigma^{\text{SGD}}(x):=\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)^{T}\right]$

######  Theorem A.16 (Stochastic modified equations).

Let $0<\eta<1,T>0$ and set $N=\lfloor T/\eta\rfloor$. Let
$x_{k}\in\mathbb{R}^{d},0\leq k\leq N$ denote a sequence of SAM iterations
defined by Eq. (4). Additionally, let us take
$\rho=\mathcal{O}\left(\eta^{1}\right).$ (107) Consider the stochastic process
$X_{t}$ defined in Eq. (106) and fix some test function $g\in G$ and suppose
that $g$ and its partial derivatives up to order 6 belong to $G$. Then, under
Assumption A.3, there exists a constant $C>0$ independent of $\eta$ such that
for all $k=0,1,\ldots,N$, we have
$\left|\mathbb{E}g\left(X_{k\eta}\right)-\mathbb{E}g\left(x_{k}\right)\right|\leq
C\eta^{1}.$ That is, the SDE (106) is an order $1$ weak approximation of the
SAM iterations (4).

######  Lemma A.17.

Under the assumptions of Theorem A.16, let $0<\eta<1$. Consider $x_{k},k\geq
0$ satisfying the SAM iterations

$x_{k+1}=x_{k}-\eta\nabla f_{\gamma_{k}}\left(x_{k}+\rho\frac{\nabla
f_{\gamma_{k}}(x_{k})}{\lVert\nabla f_{\gamma_{k}}(x_{k})\rVert}\right)$

with $x_{0}=x\in\mathbb{R}^{d}$. From the definition the one-step difference
$\bar{\Delta}=x_{1}-x$, then we have

  1. 1.

$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}f(x)\eta+\mathcal{O}(\eta^{2})\quad\forall
i=1,\ldots,d$.

  2. 2.

$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)\quad\forall
i,j=1,\ldots,d$.

  3. 3.

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{3}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}.$

All functions above are evaluated at $x$.

###### Proof of Lemma A.17.

First of all, we write that

$\partial_{e_{i}}f_{\gamma}\left(x+\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}\right)=\partial_{e_{i}}f_{\gamma}(x)+\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,0}\left(\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert}\right),$ (108)

where the residual is defined in Eq. (4) of (Folland, 2005). Therefore, for
some constant $c\in(0,1)$, it holds that

$\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,0}\left(\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}\right)=\sum_{|\alpha|=1}\frac{\partial_{e_{i}+\alpha}^{2}f_{\gamma}\left(x+c\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}\right)\rho^{1}\left(\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert}\right)^{\alpha}}{\alpha!}.$
(109)

Let us now observe that
$\mathcal{R}^{\partial_{e_{i}}f_{\gamma}(x)}_{x,0}\left(\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla f_{\gamma}(x)\rVert}\right)$ is a finite sum of
products of functions in $G$ and that, therefore, it lies in $G$.
Additionally, given its expression Eq. (109), we can factor out a common
$\rho$ and have that $K(x)=\rho K_{1}(x)$ for some function $K_{1}(x)\in G$.
Therefore, we rewrite Eq. (108) as

$\partial_{e_{i}}f_{\gamma}\left(x+\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}\right)=\partial_{e_{i}}f_{\gamma}(x)+\rho K_{1}(x).$
(110)

First of all, we notice that if we define
$\bar{K}_{1}(x)=\mathbb{E}\left[K_{1}(x)\right]$, also $\bar{K}_{1}(x)\in G$.
Therefore, it holds that

$\mathbb{E}\left[\partial_{e_{i}}f_{\gamma}\left(x+\rho\frac{\nabla
f_{\gamma}(x)}{\lVert\nabla
f_{\gamma}(x)\rVert}\right)\right]\overset{\eqref{eq:SAM_SGD_Estim_Rewritten}}{=}\partial_{e_{i}}f(x)+\rho\bar{K}_{1}(x)$
(111)

Therefore, using assumption (107), $\forall i=1,\ldots,d$, we have that

$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}f(x)\eta+\eta\rho\bar{K}_{i}(x)=-\partial_{e_{i}}f(x)\eta+\mathcal{O}\left(\eta^{2}\right)$
(112)

Additionally, by keeping in mind the definition of the covariance matrix
$\Sigma$, We immediately have

$\displaystyle\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\overset{\eqref{eq:SAM_SGD_Estim_Rewritten}}{=}$
$\displaystyle
Cov(\bar{\Delta}_{i},\bar{\Delta}_{j})+\mathbb{E}\bar{\Delta}_{i}\mathbb{E}\bar{\Delta}_{j}$
$\displaystyle=\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\eta^{2}\rho(\partial_{e_{i}}f\bar{K}_{j}(x)+\partial_{e_{j}}f\bar{K}_{i}(x))+\eta^{2}\rho^{2}\bar{K}_{i}(x)\bar{K}_{j}(x)$
$\displaystyle=\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\mathcal{O}\left(\eta^{2}\rho\right)+\mathcal{O}\left(\eta^{2}\rho^{2}\right)$
$\displaystyle=\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)\quad\forall
i,j=1,\ldots,d$ (113)

Finally, with analogous considerations, it is obvious that under our
assumptions

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{3}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}.$

∎

###### Proof of Theorem A.16.

To prove this result, all we need to do is check the conditions in Theorem
A.2. As we apply Lemma A.1, we make the following choices:

  * •

$b(x)=-\nabla f\left(x\right)$,

  * •

$\sigma(x)=\Sigma^{\text{SGD}}(X_{t})^{\frac{1}{2}}$;

First of all, we notice that $\forall i=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\overset{\text{1. Lemma
\ref{lemma:SAM_SGD}}}{=}-\partial_{e_{i}}f(x)\eta+\mathcal{O}\left(\eta^{2}\right)$;

  * •

$\mathbb{E}\Delta_{i}\overset{\text{1. Lemma
\ref{lemma:li1}}}{=}-\partial_{e_{i}}f(x)\eta+\mathcal{O}\left(\eta^{2}\right)$.

Therefore, we have that for some $K_{1}(x)\in G$

$\left|\mathbb{E}\Delta_{i}-\mathbb{E}\bar{\Delta}_{i}\right|\leq
K_{1}(x)\eta^{2},\quad\forall i=1,\ldots,d.$ (114)

Additionally,we notice that $\forall i,j=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\overset{\text{2. Lemma
\ref{lemma:SAM_SGD}}}{=}\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\Delta_{i}\Delta_{j}\overset{\text{2. Lemma
\ref{lemma:li1}}}{=}\partial_{e_{i}}f\partial_{e_{j}}f\eta^{2}+\Sigma^{\text{SGD}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{2}(x)\in G$

$\left|\mathbb{E}\Delta_{i}\Delta_{j}-\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\right|\leq
K_{2}(x)\eta^{2},\quad\forall i,j=1,\ldots,d$ (115)

Additionally, we notice that $\forall s\geq 3,\forall
i_{j}\in\\{1,\ldots,d\\}$, it holds that

  * •

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:SAM_SGD}}}{=}\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:li1}}}{=}\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{3}(x)\in G$

$\left|\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}-\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\right|\leq
K_{3}(x)\eta^{2}.$ (116)

Additionally, for some $K_{4}(x)\in G$, $\forall i_{j}\in\\{1,\ldots,d\\}$

$\mathbb{E}\prod_{j=1}^{3}\left|\bar{\Delta}_{\left(i_{j}\right)}\right|\overset{\text{3.
Lemma \ref{lemma:SAM_SGD}}}{\leq}K_{4}(x)\eta^{2}.$ (117)

Finally, Eq. (114), Eq. (115), Eq. (116), and Eq. (117) allow us to conclude
the proof.

∎

##  Appendix B Random SAM

Following (Ujváry et al., 2022) (Algorithm 2), we define Random SAM (RSAM) as
the following discrete algorithm

$x_{k+1}=x_{k}-\eta\mathbb{E}_{\epsilon\sim\mathcal{N}(0,\Sigma)}\nabla
f\gamma_{k}\left(x_{k}+\epsilon\right).$ (118)

As a first attempt, we focus on the case where $\Sigma=\sigma^{2}I_{d}$.

###  B.1 Formal Derivation - RSAM

We will consider the stochastic process $X_{t}\in\mathbb{R}^{d}$ defined by

$dX_{t}=-\nabla\tilde{f}^{\text{RSAM}}(X_{t})dt+\sqrt{\eta}\left(\Sigma^{\text{SGD}}(X_{t})+\frac{\sigma^{2}}{2}\left(\tilde{\Sigma}(X_{t})+\tilde{\Sigma}(X_{t})^{\top}\right)\right)^{\frac{1}{2}}dW_{t}$
(119)

where

$\Sigma^{\text{SGD}}(x):=\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)^{T}\right]$

is the usual covariance of SGD, while

$\tilde{\Sigma}(x):=\mathbb{E}\left[\left(\nabla f\left(x\right)-\nabla
f_{\gamma}\left(x\right)\right)\left(\mathbb{E}\left[\nabla^{3}f_{\gamma}(x)[I_{d}]\right]-\nabla^{3}f_{\gamma}(x)[I_{d}]\right)^{\top}\right]$
(120)

and

$\tilde{f}^{\text{RSAM}}(x):=f(x)+\frac{\sigma^{2}}{2}Tr(\nabla^{2}f(x)).$

In the following, we will use the notation

$\Sigma^{\text{RSAM}}(x):=\left(\Sigma^{\text{SGD}}(X_{t})+\frac{\sigma^{2}}{2}\left(\tilde{\Sigma}(X_{t})+\tilde{\Sigma}(X_{t})^{\top}\right)\right).$
(121)

######  Theorem B.1 (Stochastic modified equations).

Let $0<\eta<1,T>0$ and set $N=\lfloor T/\eta\rfloor$. Let
$x_{k}\in\mathbb{R}^{d},0\leq k\leq N$ denote a sequence of RSAM iterations
defined by Eq. (118). Additionally, let us take
$\sigma=\mathcal{O}\left(\eta^{\frac{1}{3}}\right).$ (122) Consider the
stochastic process $X_{t}$ defined in Eq. (119) and fix some test function
$g\in G$ and suppose that $g$ and its partial derivatives up to order 6 belong
to $G$. Then, under Assumption A.3, there exists a constant $C>0$ independent
of $\eta$ such that for all $k=0,1,\ldots,N$, we have
$\left|\mathbb{E}g\left(X_{k\eta}\right)-\mathbb{E}g\left(x_{k}\right)\right|\leq
C\eta^{1}.$ That is, the SDE (119) is an order $1$ weak approximation of the
RSAM iterations (118).

######  Lemma B.2.

Under the assumptions of Theorem B.1, let $0<\eta<1$ and consider $x_{k},k\geq
0$ satisfying the RSAM iterations (118)
$x_{k+1}=x_{k}-\eta\mathbb{E}_{\epsilon\sim\mathcal{N}(0,\Sigma)}\nabla
f\gamma_{k}\left(x_{k}+\epsilon\right).$ where $\Sigma=\sigma^{2}I_{d}$ and
$x_{0}=x\in\mathbb{R}^{d}$. From the definition the one-step difference
$\bar{\Delta}=x_{1}-x$, then we have 1.
$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}\tilde{f}^{\text{RSAM}}(x)\eta+\mathcal{O}(\eta\sigma^{3})\quad\forall
i=1,\ldots,d$. 2.
$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=\partial_{e_{i}}\tilde{f}^{\text{RSAM}}(x)\partial_{e_{j}}\tilde{f}^{\text{RSAM}}(x)\eta^{2}+\Sigma^{\text{RSAM}}_{(ij)}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)\quad\forall
i,j=1,\ldots,d$. 3.
$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{3}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}.$ and all the functions above are
evaluated at $x$.

###### Proof of Lemma B.2.

We perform a Taylor expansion of $\partial_{i}f(\cdot)$ around $x_{k}$

$\displaystyle
x_{k+1}^{i}=x_{k}^{i}-\mathbb{E}_{\epsilon\sim\mathcal{N}(0,\Sigma)}\left[\eta\partial_{i}f\left(x_{k}\right)-\eta\sum_{j}\partial_{ij}^{2}f\left(x_{k}\right)\epsilon_{k}^{j}-\frac{\eta}{2}\sum_{j,l}\partial_{ijl}^{3}f\left(x_{k}\right)\epsilon_{k}^{j}\epsilon_{k}^{l}+\mathcal{O}\left(\eta\left\|\epsilon_{k}\right\|^{3}\right)\right],$

and we notice that the term
$\frac{\eta}{2}\sum_{j,l}\partial_{ijl}^{3}f\left(x_{k}\right)\epsilon_{k}^{j}\epsilon_{k}^{l}$
is equal to
$\frac{\eta}{2}\partial_{i}\sum_{jl}\partial_{jl}^{2}f\left(x_{k}\right)\epsilon_{k}^{j}\epsilon_{k}^{l}$
due to Clairaut’s theorem (assuming that $f$ has continuous fourth-order
partial derivatives). By exploiting that $\epsilon_{k}$ has mean zero and
covariance $\sigma^{2}I_{d}$, we have that

$\mathbb{E}\left[x_{k+1}-x_{k}\right]=-\eta\nabla\tilde{f}^{\text{RSAM}}\left(x_{k}\right)+\mathcal{O}\left(\eta\mathbb{E}\left[\left\|\epsilon_{k}\right\|^{3}\right]\right)=-\eta\nabla\tilde{f}^{\text{RSAM}}\left(x_{k}\right)+\mathcal{O}\left(\eta\sigma^{3}\right),$

where the modified loss $\tilde{f}^{\text{RSAM}}$ is given by

$\tilde{f}^{\text{RSAM}}(x):=f(z)+\frac{\sigma^{2}}{2}\operatorname{Tr}\left(\nabla^{2}f(x)\right).$

Therefore, using (122), we have that $\forall i=1,\ldots,d$

$\mathbb{E}\bar{\Delta}_{i}=-\partial_{e_{i}}\tilde{f}^{\text{RSAM}}(x)\eta+\mathcal{O}\left(\eta^{2}\right).$
(123)

Additionally, we have that

$\displaystyle\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=\partial_{e_{i}}\tilde{f}^{\text{RSAM}}\partial_{e_{j}}\tilde{f}^{\text{RSAM}}\eta^{2}+\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})+\mathcal{O}\left(\eta^{3}\right)\quad\forall
i,j=1,\ldots,d.$ (124)

Let us now recall the expression (120) of $\tilde{\Sigma}$ and the expression
(121) of $\Sigma^{\text{RSAM}}$. Then, we automatically have that

$\text{Cov}(\bar{\Delta}_{i},\bar{\Delta}_{j})=\eta^{2}\left(\Sigma^{\text{SGD}}_{i,j}(x)+\frac{\sigma^{2}}{2}\left[\tilde{\Sigma}_{i,j}(x)+\tilde{\Sigma}_{i,j}(x)^{\top}\right]+\mathcal{O}(\sigma^{3})\right)=\eta^{2}\Sigma^{\text{RSAM}}_{i,j}(x)+\mathcal{O}(\eta^{2}\sigma^{3})$
(125)

Therefore, remembering Eq. (124) and Eq. (122) we have

$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}=\partial_{e_{i}}\tilde{f}^{\text{RSAM}}\partial_{e_{j}}\tilde{f}^{\text{RSAM}}\eta^{2}+\Sigma^{\text{SAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right),\quad\forall
i,j=1,\ldots,d$ (126)

Finally, with analogous considerations, it is obvious that under our
assumptions

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{s}\right)\quad\forall
s\geq 3,\quad i_{j}\in\\{1,\ldots,d\\}$

which in particular implies that

$\mathbb{E}\prod_{j=1}^{3}\bar{\Delta}_{i_{j}}=\mathcal{O}\left(\eta^{3}\right),\quad
i_{j}\in\\{1,\ldots,d\\}.$

∎

###### Proof of Theorem B.1.

To prove this result, all we need to do is check the conditions in Theorem
A.2. As we apply Lemma A.1, we make the following choices:

  * •

$b(x)=-\nabla\tilde{f}^{\text{RSAM}}\left(x\right)$;

  * •

$\sigma(x)=\Sigma^{\text{RSAM}}(x)^{\frac{1}{2}}$.

First of all, we notice that $\forall i=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\overset{\text{1. Lemma
\ref{lemma:RSAM_SDE}}}{=}-\partial_{e_{i}}\tilde{f}^{\text{RSAM}}(x)\eta+\mathcal{O}(\eta^{2})$;

  * •

$\mathbb{E}\Delta_{i}\overset{\text{1. Lemma
\ref{lemma:li1}}}{=}-\partial_{e_{i}}\tilde{f}^{\text{RSAM}}(x)\eta+\mathcal{O}\left(\eta^{2}\right)$.

Therefore, we have that for some $K_{1}(x)\in G$

$\left|\mathbb{E}\Delta_{i}-\mathbb{E}\bar{\Delta}_{i}\right|\leq
K_{1}(x)\eta^{2},\quad\forall i=1,\ldots,d.$ (127)

Additionally,we notice that $\forall i,j=1,\ldots,d$, it holds that

  * •

$\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\overset{\text{2. Lemma
\ref{lemma:RSAM_SDE}}}{=}\partial_{e_{i}}\tilde{f}^{\text{RSAM}}\partial_{e_{j}}\tilde{f}^{\text{RSAM}}\eta^{2}+\Sigma^{\text{RSAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\Delta_{i}\Delta_{j}\overset{\text{2. Lemma
\ref{lemma:li1}}}{=}\partial_{e_{i}}\tilde{f}^{\text{RSAM}}\partial_{e_{j}}\tilde{f}^{\text{RSAM}}\eta^{2}+\Sigma^{\text{RSAM}}_{i,j}\eta^{2}+\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{2}(x)\in G$

$\left|\mathbb{E}\Delta_{i}\Delta_{j}-\mathbb{E}\bar{\Delta}_{i}\bar{\Delta}_{j}\right|\leq
K_{2}(x)\eta^{2},\quad\forall i,j=1,\ldots,d$ (128)

Additionally, we notice that $\forall s\geq 3,\forall
i_{j}\in\\{1,\ldots,d\\}$, it holds that

  * •

$\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:RSAM_SDE}}}{=}\mathcal{O}\left(\eta^{3}\right)$;

  * •

$\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}\overset{\text{3. Lemma
\ref{lemma:li1}}}{=}\mathcal{O}\left(\eta^{3}\right)$.

Therefore, we have that for some $K_{3}(x)\in G$

$\left|\mathbb{E}\prod_{j=1}^{s}\Delta_{i_{j}}-\mathbb{E}\prod_{j=1}^{s}\bar{\Delta}_{i_{j}}\right|\leq
K_{3}(x)\eta^{2}.$ (129)

Additionally, for some $K_{4}(x)\in G$, $\forall i_{j}\in\\{1,\ldots,d\\}$

$\mathbb{E}\prod_{j=1}^{3}\left|\bar{\Delta}_{\left(i_{j}\right)}\right|\overset{\text{3.
Lemma \ref{lemma:RSAM_SDE}}}{\leq}K_{4}(x)\eta^{2}.$ (130)

Finally, Eq. (127), Eq. (128), Eq. (129), and Eq. (130) allow us to conclude
the proof.

∎

##  Appendix C Convergence Analysis: Quadratic Loss

###  C.1 ODE USAM

Let us study the quadratic loss function $f(x)=x^{\top}Hx$ where $H$ is a
diagonal matrix of eigenvalues $(\lambda_{1},\dots,\lambda_{d})$ such that
$\lambda_{1}\geq\lambda_{1}\geq\dots\geq\lambda_{d}$. Under the dynamics of
the ODE of USAM, we have that

$dX_{t}=-H\left(X_{t}+\rho HX_{t}\right)dt=-H\left(I_{d}+\rho
H\right)X_{t}dt,$ (131)

which, for the single component gives us the following dynamics

$dX^{j}_{t}=-\lambda_{j}(1+\rho\lambda_{j})X^{j}_{t}dt$ (132)

whose solution is

$X^{j}_{t}=X^{j}_{0}e^{-\lambda_{j}(1+\rho\lambda_{j})t}.$ (133)

######  Lemma C.1.

For all $\rho>0$, if all the eigenvalues of $H$ are positive, then
$X^{j}_{t}\overset{t\rightarrow\infty}{\rightarrow}0,\quad\forall
j\in\\{1,\dots,d\\}$ (134)

###### Proof of Lemma C.1.

For each $j\in\\{1,\cdots,d\\}$, we have that

$X^{j}_{t}=X^{j}_{0}e^{-\lambda_{j}(1+\rho\lambda_{j})}.$

Therefore, since the exponent is always negative, $X^{j}_{t}\rightarrow 0$ as
$t\rightarrow\infty$.

∎

######  Lemma C.2.

Let $H$ have at least one strictly negative eigenvalue and let $\lambda_{*}$
be the largest negative eigenvalue of $H$. Then, for all
$\rho>-\frac{1}{\lambda_{*}}$,
$X^{j}_{t}\overset{t\rightarrow\infty}{\rightarrow}0,\quad\forall
j\in\\{1,\dots,d\\}.$ (135)

###### Proof of Lemma C.2.

For each $j\in\\{1,\cdots,d\\}$, we have that

$X^{j}_{t}=X^{j}_{0}e^{-\lambda_{j}(1+\rho\lambda_{j})}.$

Therefore, if $\lambda_{j}>0$, the exponent is always negative for each value
of $\rho>0$. Therefore, $X^{j}_{t}\rightarrow 0$ as $t\rightarrow\infty$.
Differently, if $\lambda_{j}<0$, the exponent
$-\lambda_{j}(1+\rho\lambda_{j})$ is negative only if
$\rho>-\frac{1}{\lambda_{*}}$ where $\lambda_{*}$ is the largest negative
eigenvalue of $H$. Therefore, if$\rho>-\frac{1}{\lambda_{*}}$,
$X^{j}_{0}\rightarrow 0$ if $t\rightarrow\infty$.

∎

###  C.2 SDE USAM - Stationary Distribution

Let us consider the noisy quadratic model $f(x)=\frac{1}{2}x^{\top}Hx$, where
$H$ is a symmetric matrix. Then, based on Theorem (A.4) in the case where
$\Sigma(x)=\varsigma I_{d}$, the corresponding SDE is give by

$dX_{t}=-H\left(I_{d}+\rho H\right)X_{t}dt+\left[(I_{d}+\rho
H)\sqrt{\eta}\varsigma\right]dW_{t}.$ (136)

######  Theorem C.3 (Stationary distribution - PSD Case.).

For any $\rho>0$, the stationary distribution of Eq. (136) is $\displaystyle
P\left(x,\infty\mid\rho\right)=\sqrt{\frac{\lambda_{i}}{\pi\eta\varsigma^{2}}\frac{1}{1+\rho\lambda_{i}}}\exp\left[-\frac{\lambda_{i}}{\eta\varsigma^{2}}\frac{1}{1+\rho\lambda_{i}}x^{2}\right]$
(137) where $(\lambda_{1},\dots,\lambda_{d})$ are the eigenvalues of $H$ and
$\lambda_{i}>0,\forall i\in\\{1,\cdots,d\\}$.

More interestingly, if $\rho$ is too large, this same conclusion holds even
for a saddle point.

######  Theorem C.4 (Stationary distribution - Indefinite Case.).