SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Cases A-1 and A-2 show that  $x_{t+1}$  also remains in  $[c - \rho, c + \rho]$ , when we update  $m$ -SAM using  $G_{f^{(1)}}(x_t)$ .
Figure 4: The original and virtual loss plot for the example function in Theorem 4.2. The graph drawn in purple and red are the original/virtual loss of component functions. The graph drawn in green indicates  $f$ , and the graph drawn in blue indicates  $\mathbb{E}J_{f}$ . (a)  $f$  and its component functions  $f^{(1)}$ ,  $f^{(2)}$ . (b)  $\mathbb{E}J_{f}$  and its component functions  $J_{f^{(1)}}$ ,  $J_{f^{(2)}}$ .
We next examine the case involving  $G_{f^{(2)}}$ . The next iterate  $x_{t + 1}$  of  $m$ -SAM using  $G_{f^{(2)}}(x_t)$  is
$$
x _ {t + 1} = \left\{ \begin{array}{l l} x _ {t} - \eta \frac {a}{1 - p} \left(x _ {t} - \rho + p (x _ {t} - \rho - c)\right), & x _ {t} <   c ^ {\prime} \\ x _ {t}, & x _ {t} = c ^ {\prime} \\ x _ {t} - \eta \frac {a}{1 - p} \left(x _ {t} + \rho + p (x _ {t} + \rho - c)\right), & x _ {t} > c ^ {\prime}. \end{array} \right.
$$
Again, we divide it into two cases, since the  $x_{t} = c^{\prime}$  case is obvious.
Case B-1:  $c - \rho \leq x_{t} < c'$ . Since  $(1 + p)x_{t} - (1 + p)\rho - pc < 0$  for  $x_{t} \in [c - \rho, c')$ , we have
$$
\begin{array}{l} c - \rho \leq x _ {t} \leq x _ {t + 1} = x _ {t} - \eta \frac {a}{1 - p} ((1 + p) x _ {t} - (1 + p) \rho - p c) \\ = \left(1 - \eta \frac {a (1 + p)}{1 - p}\right) x _ {t} + \eta \frac {a (1 + p)}{1 - p} \rho + \eta \frac {a p}{1 - p} c. \\ \end{array}
$$
Since  $\eta a \leq \frac{3}{50}$ , we have  $1 - \eta \frac{a(1 + p)}{1 - p} = 1 - 5\eta a \geq 0$ , so
$$
\begin{array}{l} \left(1 - \eta \frac {a (1 + p)}{1 - p}\right) x _ {t} + \eta \frac {a (1 + p)}{1 - p} \rho + \eta \frac {a p}{1 - p} c \\ \leq \left(1 - \eta \frac {a (1 + p)}{1 - p}\right) c ^ {\prime} + \eta \frac {a (1 + p)}{1 - p} \rho + \eta \frac {a p}{1 - p} c \\ = c ^ {\prime} + \eta \frac {a (1 + p)}{1 - p} \rho = c ^ {\prime} + 5 \eta a \rho \leq c ^ {\prime} + \rho \leq c + \rho . \\ \end{array}
$$
Case B-2:  $c' < x_t \leq c + \rho$ . Since  $(1 + p)x_t + (1 + p)\rho - pc > 0$  for  $x_t \in (c', c + \rho]$ , we have
$$
\begin{array}{l} c + \rho \geq x _ {t} \geq x _ {t + 1} = x _ {t} - \eta \frac {a}{1 - p} ((1 + p) x _ {t} + (1 + p) \rho - p c) \\ \geq c ^ {\prime} - \eta \frac {a}{1 - p} ((1 + p) c ^ {\prime} + (1 + p) \rho - p c) \\ = \frac {p}{1 + p} c - \eta \frac {a (1 + p)}{1 - p} \rho = c - \frac {1}{1 + p} c - 5 \eta a \rho \\ = c - \frac {1}{1 + p} \cdot \left(1 + \frac {p}{4}\right) \rho - 5 \eta a \rho \\ \geq c - \frac {3}{5} \cdot \left(\frac {7}{6}\right) \rho - \frac {3}{1 0} \rho = c - \rho , \\ \end{array}
$$
where the last inequality used  $\eta a\leq \frac{3}{50}$
Cases B-1 and B-2 demonstrate that  $x_{t + 1}$  also remains within the interval  $[c - \rho, c + \rho]$  when we update  $m$ -SAM using  $G_{f^{(2)}}(x_t)$ .
To gain a better intuitive understanding, we provide Figures 4(a), 4(b), demonstrating the original and virtual loss functions of  $f^{(1)}$  and  $f^{(2)}$ . In Figure 4(b), we can examine that  $m$ -SAM updates generate an attraction basin in  $J_{f^{(1)}}$ , thereby making a region that  $m$ -SAM cannot escape.
The aforementioned case analyses indicate that when the initial point  $x_0$  falls within the interval  $[c - \rho, c + \rho]$ , all subsequent iterations of  $m$ -SAM will also remain within this interval, regardless of the selected component function for updating. Given  $c = \left(1 + \frac{p}{4}\right)\rho = \frac{7}{6}\rho$ , we can conclude that all subsequent iterations of  $m$ -SAM are points located at a distance of at least  $\frac{1}{6}\rho$  from the global optimum, for any  $\beta > 0$ ,  $\rho > 0$ , and  $\sigma > 0$ .
Moreover, in case of  $\sigma >\beta \rho$  , it follows that  $f(x) = \frac{a}{2} x^{2} = \frac{\beta}{10} x^{2}$  . Consequently, the suboptimality gap at any timestep is at least
$$
f (x _ {t}) - f ^ {*} = f (x _ {t}) \geq f \left(\frac {\rho}{6}\right) = \Omega (\beta \rho^ {2}),
$$
C.4 Convergence Proof for Smooth and Convex Functions (Proof of Theorem 4.3)
In this section, we establish the convergence of stochastic SAM for smooth and convex functions to near-stationary points. For ease of understanding, we provide the theorem statement here.
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| \nabla f (\pmb {x} _ {t}) \| ^ {2} = \mathcal {O} \bigg (\frac {\beta \Delta}{T} + \frac {\sqrt {\beta [ \sigma^ {2} - \beta^ {2} \rho^ {2} ] _ {+} \Delta}}{\sqrt {T}} \bigg) + 4 \beta^ {2} \rho^ {2}.
$$
Proof. We start from Lemma C.3 with  $\mu = 0$ . In order to do this, additionally assuming  $\beta$ -smoothness of component functions  $l(\cdot, \xi)$  is necessary for  $m$ -SAM.
$$
\mathbb {E} f (\boldsymbol {x} _ {t + 1}) \leq \mathbb {E} f (\boldsymbol {x} _ {t}) - \frac {\eta}{2} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} + 2 \beta^ {2} \rho^ {2} \eta + \beta \eta^ {2} (\sigma^ {2} - \beta^ {2} \rho^ {2}),
$$
$$
\mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} \leq \frac {2}{\eta} \left(\mathbb {E} f (\boldsymbol {x} _ {t}) - \mathbb {E} f (\boldsymbol {x} _ {t + 1})\right) + 4 \beta^ {2} \rho^ {2} + 2 \beta \eta \left(\sigma^ {2} - \beta^ {2} \rho^ {2}\right).
$$
Adding up the inequality for  $t = 0, \dots, T - 1$ , and dividing both sides by  $T$ , we get
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} \leq \frac {2}{\eta T} \left(\mathbb {E} f (\boldsymbol {x} _ {0}) - \mathbb {E} f (\boldsymbol {x} _ {T})\right) + 4 \beta^ {2} \rho^ {2} + 2 \beta \eta \left(\sigma^ {2} - \beta^ {2} \rho^ {2}\right) \\ \leq \frac {2}{\eta T} \Delta + 4 \beta^ {2} \rho^ {2} + 2 \eta \beta \left(\sigma^ {2} - \beta^ {2} \rho^ {2}\right). \tag {11} \\ \end{array}
$$
The convergence rate varies depending on two different cases determined by the value of  $\sigma$ .
Case A:  $\sigma \leq \beta \rho$ . In this case, it must hold that  $\eta = \frac{1}{2\beta}$ . By this choice of  $\eta$ , we have
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} \leq \frac {2 \Delta}{\eta T} + 4 \beta^ {2} \rho^ {2} \\ = \frac {4 \beta \Delta}{T} + 4 \beta^ {2} \rho^ {2}. \\ \end{array}
$$
Case B:  $\sigma >\beta \rho$  Setting  $\eta = \min \left\{\sqrt{\frac{\Delta}{\beta(\sigma^2 - \beta^2\rho^2)T}},\frac{1}{2\beta}\right\}$ , we consider two cases.
Case B-1:  $\sqrt{\frac{\Delta}{\beta(\sigma^2 - \beta^2\rho^2)T}}\leq \frac{1}{2\beta}$ . Putting  $\eta = \sqrt{\frac{\Delta}{\beta(\sigma^2 - \beta^2\rho^2)T}}$  into (11), we get
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} \leq 4 \sqrt {\frac {\beta (\sigma^ {2} - \beta^ {2} \rho^ {2}) \Delta}{T}} + 4 \beta^ {2} \rho^ {2}.
$$
Case B-2:  $\sqrt{\frac{\Delta}{\beta(\sigma^2 - \beta^2\rho^2)T}}\geq \frac{1}{2\beta}$ . Placing  $\eta = \frac{1}{2\beta}$  into (11), we get
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} \leq \frac {4 \beta \Delta}{T} + 4 \beta^ {2} \rho^ {2} + 2 \cdot \frac {1}{2 \beta} \cdot \beta (\sigma^ {2} - \beta^ {2} \rho^ {2}) \\ \leq \frac {4 \beta \Delta}{T} + 4 \beta^ {2} \rho^ {2} + 2 \cdot \sqrt {\frac {\Delta}{\beta (\sigma^ {2} - \beta^ {2} \rho^ {2}) T}} \cdot \beta (\sigma^ {2} - \beta^ {2} \rho^ {2}) \\ = \frac {4 \beta \Delta}{T} + 2 \sqrt {\frac {\beta (\sigma^ {2} - \beta^ {2} \rho^ {2}) \Delta}{T}} + 4 \beta^ {2} \rho^ {2}. \\ \end{array}
$$
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| \nabla f (\pmb {x} _ {t}) \| ^ {2} = \mathcal {O} \left(\frac {\beta \Delta}{T} + \frac {\sqrt {\beta (\sigma^ {2} - \beta^ {2} \rho^ {2}) \Delta}}{\sqrt {T}}\right) + 4 \beta^ {2} \rho^ {2},
$$
C.5 Non-Convergence of  $m$ -SAM for a Smooth and Convex Function (Proof of Theorem 4.4)
In this section, we provide an in-depth analysis of our counterexample, which demonstrates that stochastic  $m$ -SAM with constant perturbation size  $\rho$  can provably converge to a point that is not a global minimum. Furthermore, the suboptimality gap in terms of the function value can be made arbitrarily large, hence indicating that proving convergence of  $m$ -SAM to global minima (modulo some additive factors  $\mathcal{O}(\rho^2)$ ) is impossible. The theorem is restated for convenience.
Proof. Given  $\rho > 0$ ,  $\beta > 0$ ,  $\sigma > 0$ , we set an arbitrary constant  $c > \frac{5\rho}{4}$ , and a parameter  $a > 0$  which will be chosen later. For  $x \in \mathbb{R}$ , consider a one-dimensional smooth convex function,
$$
f (x) = \left\{ \begin{array}{l l} a x + \frac {\beta}{2} x ^ {2}, & x \leq 0 \\ a x, & x \geq 0, \end{array} \right.
$$
and  $l(x;\xi)$  can be given as
where the functions  $f^{(1)}$  and  $f^{(2)}$  are given by the following definitions.
$$
f ^ {(1)} (x) = \left\{ \begin{array}{l l} 2 a \left(x - c + \frac {\rho}{8}\right) + \frac {\beta}{2} x ^ {2}, & x \leq 0 \\ 2 a \left(x - c + \frac {\rho}{8}\right), & 0 \leq x \leq c - \frac {\rho}{4} \\ - \frac {4 a}{\rho} \left(x - c\right) ^ {2}, & c - \frac {\rho}{4} \leq x \leq c + \frac {\rho}{4} \\ - 2 a \left(x - c - \frac {\rho}{8}\right), & c + \frac {\rho}{4} \leq x, \end{array} \right.
$$
$$
f ^ {(2)} (x) = \left\{ \begin{array}{l l} \frac {1}{1 - p} \left(a x - 2 p a \left(x - c + \frac {\rho}{8}\right)\right) + \frac {\beta}{2} x ^ {2}, & x \leq 0 \\ \frac {1}{1 - p} \left(a x - 2 p a \left(x - c + \frac {\rho}{8}\right)\right), & 0 \leq x \leq c - \frac {\rho}{4} \\ \frac {1}{1 - p} \left(a x + \frac {4 p a}{\rho} (x - c) ^ {2}\right), & c - \frac {\rho}{4} \leq x \leq c + \frac {\rho}{4} \\ \frac {1}{1 - p} \left(a x + 2 p a \left(x - c - \frac {\rho}{8}\right)\right), & c + \frac {\rho}{4} \leq x. \end{array} \right.
$$
It can be verified that  $f^{(1)}$  is  $\left(\max \left\{\frac{8a}{\rho}, \beta\right\}\right)$ -smooth, and  $f^{(2)}$  is  $\left(\max \left\{\frac{8a}{\rho} \cdot \frac{p}{1 - p}, \beta\right\}\right)$ -smooth. Moreover,  $\| \nabla f(x) - \nabla f^{(1)}(x) \|^2 \leq 9a^2$  holds, as well as  $\| \nabla f(x) - \nabla f^{(2)}(x) \|^2 \leq 9a^2 \cdot \frac{p^2}{(1 - p)^2}$ . Consequently,  $\mathbb{E} \| \nabla f(x) - g(x) \|^2 \leq 9a^2 \cdot \frac{p}{1 - p}$ .
By selecting  $p > \frac{1}{2}$ , and setting  $a = \min \left\{\frac{\beta\rho(1 - p)}{8p},\frac{\sigma\sqrt{1 - p}}{3\sqrt{p}}\right\}$ , we can check that  $\mathbb{E}l(x;\xi) = f(x)$  and the component functions are  $\beta$ -smooth. Additionally,  $f$  satisfies Assumption 2.5.
In the following analysis, we examine the virtual gradient maps  $G_{f^{(1)}}$  and  $G_{f^{(2)}}$  in the specified region of interest:  $c - \rho \leq x \leq c + \rho$ . For this specific interval, we are going to show that if we start  $m$ -SAM inside this interval, then all the subsequent iterates of  $m$ -SAM must stay inside the same interval  $[c - \rho, c + \rho]$ .
In this region,  $y = x + \rho \frac{\nabla f^{(1)}(x)}{\|\nabla f^{(1)}(x)\|}$  and  $G_{f^{(1)}}$  are as follows.
$$
y = \left\{ \begin{array}{l l} x + \rho , & c - \rho \leq x <   c \\ x, & x = c \\ x - \rho , & c <   x \leq c + \rho . \end{array} \right.
$$
$$
G _ {f ^ {(1)}} (x) = \left\{ \begin{array}{l l} - \frac {8 a}{\rho} (x + \rho - c), & c - \rho \leq x \leq c - \frac {3 \rho}{4} \\ - 2 a, & c - \frac {3 \rho}{4} \leq x <   c \\ 0, & x = c \\ 2 a, & c <   x \leq c + \frac {3 \rho}{4} \\ - \frac {8 a}{\rho} (x - \rho - c), & c + \frac {3 \rho}{4} \leq x \leq c + \rho . \end{array} \right.
$$
Additionally defining  $c' \triangleq c - \frac{\rho}{8p}$ , we can compute  $y = x + \rho \frac{\nabla f^{(2)}(x)}{\|\nabla f^{(2)}(x)\|}$  and  $G_{f^{(2)}}$  as follows.
$$
y = \left\{ \begin{array}{l l} x - \rho , & c - \rho \leq x <   c ^ {\prime} \\ x, & x = c ^ {\prime} \\ x + \rho , & c ^ {\prime} <   x \leq c + \rho . \end{array} \right.
$$
$$
G _ {f ^ {(2)}} (x) = \left\{ \begin{array}{l l} \frac {a - 2 a p}{1 - p}, & c - \rho \leq x <   c ^ {\prime} \\ 0, & x = c ^ {\prime} \\ \frac {a + 2 a p}{1 - p}, & c ^ {\prime} <   x \leq c + \rho . \end{array} \right.
$$
Now recall that  $\eta \leq \frac{1}{\beta}$ . Given that  $a \leq \frac{\beta\rho(1 - p)}{8p}$ , we can derive  $\eta \leq \frac{1}{\beta} \leq \frac{\rho}{a} \cdot \frac{1 - p}{8p}$ . Furthermore, the next iterate  $x_{t + 1}$  of  $m$ -SAM using  $G_{f^{(1)}}(x_t)$  for  $c - \rho \leq x_t \leq c + \rho$  is
$$
x _ {t + 1} = x _ {t} - \eta G _ {f ^ {(1)}} (x _ {t}) = \left\{ \begin{array}{l l} x _ {t} + \eta \cdot \frac {8 a}{\rho} (x + \rho - c), & c - \rho \leq x _ {t} \leq c - \frac {3 \rho}{4} \\ x _ {t} + 2 \eta a, & c - \frac {3 \rho}{4} \leq x _ {t} <   c \\ x _ {t}, & x _ {t} = c \\ x _ {t} - 2 \eta a, & c <   x _ {t} \leq c + \frac {3 \rho}{4} \\ x _ {t} + \eta \cdot \frac {8 a}{\rho} (x - \rho - c), & c + \frac {3 \rho}{4} \leq x _ {t} \leq c + \rho . \end{array} \right.
$$
We will now show that  $x_{t + 1}$  must remain in the interval  $[c - \rho, c + \rho]$ . For the case  $x_t = c$ , it is obvious. Dividing the rest into three cases,
Case A-1:  $c - \rho \leq x_t \leq c - \frac{3\rho}{4}$ . Since  $0 \leq \eta \cdot \frac{8a}{\rho} (x_t + \rho - c) \leq 2\eta a \leq \rho \cdot \frac{1 - p}{4p} \leq \frac{\rho}{4}$ , we have
$$
c - \rho \leq x _ {t} \leq x _ {t + 1} \leq x _ {t} + \frac {\rho}{4} \leq c + \rho ,
$$
which proves that  $x_{t + 1}$  also remains in  $[c - \rho, c + \rho]$ .
(b)
Figure 5: The original and virtual loss plot for the example function in Theorem 4.4. The graph drawn in purple and red are the original/virtual loss of component functions. The graph drawn in green indicates  $f$ , and the graph drawn in blue indicates  $\mathbb{E}J_{f}$ . (a)  $f$  and its component functions  $f^{(1)}$ ,  $f^{(2)}$ . (b)  $\mathbb{E}J_{f}$  and its component functions  $J_{f^{(1)}}, J_{f^{(2)}}$ .
Case A-2:  $c - \frac{3\rho}{4} \leq x_t < c, c < x_t \leq c + \frac{3\rho}{4}$ . Since  $0 \leq 2\eta a \leq \rho \cdot \frac{1 - p}{4p} \leq \frac{\rho}{4}$ , we have
$$
c - \rho \leq x _ {t} - \frac {\rho}{4} \leq x _ {t + 1} \leq x _ {t} + \frac {\rho}{4} \leq c + \rho ,
$$
thereby proving that  $x_{t + 1}$  also remains in  $[c - \rho, c + \rho]$ .
Case A-3:  $c + \frac{3\rho}{4} \leq x_t \leq c + \rho$ . Since  $0 \geq \eta \cdot \frac{8a}{\rho}(x_t - \rho - c) \geq -2\eta a \geq -\rho \cdot \frac{1 - p}{4p} \geq -\frac{\rho}{4}$ , we have
$$
c - \rho \leq x _ {t} - \frac {\rho}{4} \leq x _ {t + 1} \leq x _ {t} \leq c + \rho ,
$$
which shows that  $x_{t + 1}$  also remains in  $[c - \rho, c + \rho]$ .
We next consider the case of  $G_{f^{(2)}}$ . The next iterate  $x_{t + 1}$  of  $m$ -SAM using  $G_{f^{(2)}}(x_t)$  for  $c - \rho \leq x_{t} \leq c + \rho$  is given by
$$
x _ {t + 1} = x _ {t} - \eta G _ {f ^ {(2)}} (x _ {t}) = \left\{ \begin{array}{l l} x _ {t} - \eta a \cdot \frac {1 - 2 p}{1 - p}, & c - \rho \leq x _ {t} <   c ^ {\prime} \\ x _ {t}, & x _ {t} = c ^ {\prime} \\ x _ {t} - \eta a \cdot \frac {1 + 2 p}{1 - p}, & c ^ {\prime} <   x _ {t} \leq c + \rho . \end{array} \right.
$$
Again, the  $x_{t} = c^{\prime}$  case is trivial. Considering the remaining two cases,
Case B-1:  $c - \rho \leq x_t < c'$ . Since  $0 \geq \eta a \cdot \frac{1 - 2p}{1 - p} \geq \rho \cdot \frac{1 - 2p}{8p} \geq -(c + \rho - c')$ , we get
$$
c - \rho \leq x _ {t} \leq x _ {t + 1} \leq x _ {t} + \left(c + \rho - c ^ {\prime}\right) \leq c + \rho ,
$$
which indicates that  $x_{t + 1}$  also remains in  $[c - \rho, c + \rho]$ .
Case B-2:  $c' < x_t \leq c + \rho$ . Since  $0 \leq \eta a \cdot \frac{1 + 2p}{1 - p} \leq \rho \cdot \frac{1 + 2p}{8p} \leq c' - (c - \rho)$ , we have
$$
c - \rho \leq x _ {t} - c ^ {\prime} + (c - \rho) \leq x _ {t + 1} \leq x _ {t} \leq c + \rho ,
$$
The case analyses above indicate that if  $x_0$  is initialized in  $[c - \rho, c + \rho]$ , the subsequent iterates of  $m$ -SAM also remain in  $[c - \rho, c + \rho]$ , regardless of the component function chosen by  $m$ -SAM for update. Figures 5(a) and 5(b) demonstrate the original loss function and virtual loss function of component functions  $f^{(1)}$  and  $f^{(2)}$ . Due to the concavity of  $f^{(1)}$ , a basin of attraction is formed in  $J_{f^{(1)}}$ , thereby creating a region where the iterations of  $m$ -SAM cannot escape.
Therefore, the suboptimality gap at any timestep is at least
$$
f (x _ {t}) - f ^ {*} = f (x _ {t}) - f \left(- \frac {a}{\beta}\right) \geq f (c - \rho) - f \left(- \frac {a}{\beta}\right) = a (c - \rho) + \frac {a ^ {2}}{2 \beta}.
$$
Therefore, regardless of iterate averaging scheme, the suboptimality gap in terms of function value will stay above  $a(c - \rho) + \frac{a^2}{2\beta}$ .
Moreover, this suboptimality gap can be made arbitrarily large if we choose larger values of  $c$ ; notice that  $c > \frac{5\rho}{4}$  is the only requirement on  $c$ . Consequently, it is impossible to guarantee convergence to global minima, even up to an additive factor, for  $m$ -SAM. This finishes the proof.
C.6 Convergence Proof for Smooth and Nonconvex Function under  $n$ -SAM (Proof of Theorem 4.5)
In this section, we establish the convergence result of stochastic  $n$ -SAM for smooth and nonconvex functions. For convenience, we restate the theorem.
Theorem 4.5. Consider a  $\beta$ -smooth function  $f$  satisfying  $f^{*} = \inf_{x} f(x) > -\infty$ , and assume Assumption 2.5. Under  $n$ -SAM, starting at  $\mathbf{x}_0$  with any perturbation size  $\rho > 0$  and step size  $\eta = \min \left\{\frac{1}{2\beta}, \frac{\sqrt{\Delta}}{\sqrt{\beta \sigma^2 T}}\right\}$  to minimize  $f$ , we have
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| \nabla f (\pmb {x} _ {t}) \| ^ {2} \leq \mathcal {O} \left(\frac {\beta \Delta}{T} + \frac {\sqrt {\beta \sigma^ {2} \Delta}}{\sqrt {T}}\right) + \beta^ {2} \rho^ {2}.
$$
$$
\begin{array}{l} \mathbb {E} f (\boldsymbol {x} _ {t + 1}) \leq \mathbb {E} f (\boldsymbol {x} _ {t}) - \eta \mathbb {E} \left\langle \nabla f (\boldsymbol {x} _ {t}), \tilde {g} (\boldsymbol {y} _ {t}) \right\rangle + \frac {\beta \eta^ {2}}{2} \mathbb {E} \| \tilde {g} (\boldsymbol {y} _ {t}) \| ^ {2} \\ \leq \mathbb {E} f (\boldsymbol {x} _ {t}) - \eta \mathbb {E} \left\langle \nabla f (\boldsymbol {x} _ {t}), \nabla f (\boldsymbol {y} _ {t}) \right\rangle + \beta \eta^ {2} \left(\mathbb {E} \| \nabla f (\boldsymbol {y} _ {t}) \| ^ {2} + \mathbb {E} \| \tilde {g} (\boldsymbol {y} _ {t}) - \nabla f (\boldsymbol {y} _ {t}) \| ^ {2}\right) \\ \leq \mathbb {E} f (\boldsymbol {x} _ {t}) - \eta \mathbb {E} \left\langle \nabla f (\boldsymbol {x} _ {t}), \nabla f (\boldsymbol {y} _ {t}) \right\rangle + \beta \eta^ {2} \left(\mathbb {E} \| \nabla f (\boldsymbol {y} _ {t}) \| ^ {2} + \sigma^ {2}\right) \\ = \mathbb {E} f (\boldsymbol {x} _ {t}) - \frac {\eta}{2} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} - \frac {\eta}{2} \mathbb {E} \| \nabla f (\boldsymbol {y} _ {t}) \| ^ {2} + \frac {\eta}{2} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) - \nabla f (\boldsymbol {y} _ {t}) \| ^ {2} \\ + \beta \eta^ {2} \left(\mathbb {E} \| \nabla f (\boldsymbol {y} _ {t}) \| ^ {2} + \sigma^ {2}\right) \\ \leq \mathbb {E} f (\boldsymbol {x} _ {t}) - \frac {\eta}{2} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} + \frac {\beta^ {2} \eta}{2} \mathbb {E} \| \boldsymbol {x} _ {t} - \boldsymbol {y} _ {t} \| ^ {2} + \beta \sigma^ {2} \eta^ {2} \\ = \mathbb {E} f (\boldsymbol {x} _ {t}) - \frac {\eta}{2} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} + \frac {\beta^ {2} \rho^ {2} \eta}{2} + \beta \sigma^ {2} \eta^ {2}. \\ \end{array}
$$
Rearranging the inequality, we get
$$
\mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} \leq \frac {2}{\eta} (\mathbb {E} f (\boldsymbol {x} _ {t}) - \mathbb {E} f (\boldsymbol {x} _ {t + 1})) + \beta^ {2} \rho^ {2} + 2 \beta \sigma^ {2} \eta .
$$
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} \leq \frac {2}{\eta T} \left(\mathbb {E} f (\boldsymbol {x} _ {0}) - \mathbb {E} f (\boldsymbol {x} _ {T})\right) + \beta^ {2} \rho^ {2} + 2 \beta \sigma^ {2} \eta \\ \leq \frac {2 \Delta}{\eta T} + \beta^ {2} \rho^ {2} + 2 \beta \sigma^ {2} \eta . \tag {12} \\ \end{array}
$$
Substituting  $\eta = \min \left\{\frac{1}{2\beta},\frac{\sqrt{\Delta}}{\sqrt{T\beta\sigma^2}}\right\}$  to (12), we get two cases.
Case A:  $\frac{1}{2\beta} \leq \frac{\sqrt{\Delta}}{\sqrt{T\beta\sigma^2}}$ .  $\eta = \frac{1}{2\beta}$ , so we have
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} \leq \frac {2 \Delta}{\eta T} + \beta^ {2} \rho^ {2} + 2 \beta \sigma^ {2} \eta
$$
$$
\begin{array}{l} \leq \frac {4 \beta \Delta}{T} + 2 \beta \sigma^ {2} \cdot \frac {\sqrt {\Delta}}{\sqrt {T \beta \sigma^ {2}}} + \beta^ {2} \rho^ {2} \\ = \frac {4 \beta \Delta}{T} + \frac {2 \sqrt {\beta \sigma^ {2} \Delta}}{\sqrt {T}} + \beta^ {2} \rho^ {2}. \\ \end{array}
$$
Case B:  $\frac{\sqrt{\Delta}}{\sqrt{T\beta\sigma^2}} \leq \frac{1}{2\beta}$ .  $\eta = \frac{\sqrt{\Delta}}{\sqrt{T\beta\sigma^2}}$ , and it leads to
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} \leq \frac {2 \Delta}{\eta T} + \beta^ {2} \rho^ {2} + 2 \beta \sigma^ {2} \eta \\ \leq \frac {2 \sqrt {\Delta}}{T} \cdot \sqrt {T \beta \sigma^ {2}} + \frac {2 \sqrt {\beta \sigma^ {2} \Delta}}{\sqrt {T}} + \beta^ {2} \rho^ {2} \\ = \frac {4 \sqrt {\beta \sigma^ {2} \Delta}}{\sqrt {T}} + \beta^ {2} \rho^ {2}. \\ \end{array}
$$
Merging two cases, we can conclude that
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| \nabla f (\boldsymbol {x} _ {t}) \| ^ {2} \leq \mathcal {O} \left(\frac {\beta \Delta}{T} + \frac {\sqrt {\beta \sigma^ {2} \Delta}}{\sqrt {T}}\right) + \beta^ {2} \rho^ {2},
$$
C.7 Convergence Proof for Smooth Lipschitz Nonconvex Function under  $m$ -SAM (Proof of Theorem 4.6 and Corollary 4.7)
In this section, the convergence proof for stochastic  $m$ -SAM for smooth, Lipschitz, and nonconvex functions is presented. The notation  $\hat{\pmb{y}}_t = \pmb{x}_t + \rho \frac{\nabla f(\pmb{x}_t)}{\|\nabla f(\pmb{x}_t)\|}$  is used here once again. The theorem is restated for convenience.
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \left[ \left(\| \nabla f (\boldsymbol {x} _ {t}) \| - \beta \rho\right) ^ {2} \right] \leq \mathcal {O} \left(\frac {\sqrt {\beta \Delta \left(\sigma^ {2} + L ^ {2}\right)}}{\sqrt {T}}\right) + 5 \beta^ {2} \rho^ {2}.
$$
Proof. The proof technique resembles Mi et al. [27]. Starting from the definition of  $\beta$ -smoothness, we have