Question:
You are a highly experienced, conscientious, and fair academic reviewer, please help me review this paper. The review should be organized into nine sections:
1. Summary: A summary of the paper in 100-150 words.
2. Strengths/Weaknesses/Questions: The Strengths/Weaknesses/Questions of paper, which should be listed in bullet points, with each point supported by specific examples from the article where possible.
3. Soundness/Contribution/Presentation: Rate the paper's Soundness/Contribution/Presentation, and match this score to the corresponding description from the list below and provide the result. The possible scores and their descriptions are:
    1 poor
    2 fair
    3 good
    4 excellent
4. Rating: Give this paper an appropriate rating, match this rating to the corresponding description from the list below and provide the result. The possible Ratings and their descriptions are:
    1 strong reject
    2 reject, significant issues present
    3 reject, not good enough
    4 possibly reject, but has redeeming facets
    5 marginally below the acceptance threshold
    6 marginally above the acceptance threshold
    7 accept, but needs minor improvements
    8 accept, good paper
    9 strong accept, excellent work
    10 strong accept, should be highlighted at the conference
5. Paper Decision: It must include the Decision itself(Accept or Reject) and the reasons for this decision, based on the criteria of originality, methodological soundness, significance of results, and clarity and logic of presentation.

Here is the template for a review format, you must follow this format to output your review result:
**Summary:**
Summary content

**Strengths:**
- Strength 1
- Strength 2
- ...

**Weaknesses:**
- Weakness 1
- Weakness 2
- ...

**Questions:**
- Question 1
- Question 2
- ...

**Soundness:**
Soundness result

**Presentation:**
Presentation result

**Contribution:**
Contribution result

**Rating:**
Rating result

**Paper Decision:**
- Decision: Accept/Reject
- Reasons: reasons content

Please ensure your feedback is objective and constructive. The paper is as follows:

# Why Does Sharpness-Aware Minimization Generalize Better Than SGD?

Zixiang Chen Junkai Zhang Yiwen Kou Xiangning Chen Cho-Jui Hsieh Quanquan Gu

Department of Computer Science

University of California, Los Angeles

Los Angeles, CA 90095

{chenzx19,zhang,evankou,xiangning,chohsieh,qgu}@cs.ucla.edu

Equal contribution.

###### Abstract

The challenge of overfitting, in which the model memorizes the training data and fails to generalize to test data, has become increasingly significant in the training of large neural networks. To tackle this challenge, Sharpness-Aware Minimization (SAM) has emerged as a promising training method, which can improve the generalization of neural networks even in the presence of label noise. However, a deep understanding of how SAM works, especially in the setting of nonlinear neural networks and classification tasks, remains largely missing. This paper fills this gap by demonstrating why SAM generalizes better than Stochastic Gradient Descent (SGD) for a certain data model and two-layer convolutional ReLU networks. The loss landscape of our studied problem is nonsmooth, thus current explanations for the success of SAM based on the Hessian information are insufficient. Our result explains the benefits of SAM, particularly its ability to prevent noise learning in the early stages, thereby facilitating more effective learning of features. Experiments on both synthetic and real data corroborate our theory.

## 1 Introduction

The remarkable performance of deep neural networks has sparked considerable interest in creating ever-larger deep learning models, while the training process continues to be a critical bottleneck affecting overall model performance. The training of large models is unstable and difficult due to the sharpness, non-convexity, and non-smoothness of its loss landscape. In addition, as the number of model parameters is much larger than the training sample size, the model has the ability to memorize even randomly labeled data (Zhang et al., 2021), which leads to overfitting. Therefore, although traditional gradient-based methods like gradient descent (GD) and stochastic gradient descent (SGD) can achieve generalizable models under certain conditions, these methods may suffer from unstable training and harmful overfitting in general.

To overcome the above challenge, _Sharpness-Aware Minimization_ (SAM) (Foret et al., 2020), an innovative training paradigm, has exhibited significant improvement in model generalization and has become widely adopted in many applications. In contrast to traditional gradient-based methods that primarily focus on finding a point in the parameter space with a minimal gradient norm, SAM also pursues a solution with reduced sharpness, characterized by how rapidly the loss function changes locally. Despite the empirical success of SAM across numerous tasks (Bahri et al., 2021; Behdin et al., 2022; Chen et al., 2021; Liu et al., 2022a), the theoretical understanding of this method remains limited.

Foret et al. (2020) provided a PAC-Bayes bound on the generalization error of SAM to show that it will generalize well, while the bound only holds for the infeasible average-direction perturbation instead ofpractically used ascend-direction perturbation. Andriushchenko and Flammarion (2022) investigated the implicit bias of SAM for diagonal linear networks under global convergence assumption. The oscillations in the trajectory of SAM were explored by Bartlett et al. (2022), leading to a convergence result for the convex quadratic loss. A concurrent work (Wen et al., 2022) demonstrated that SAM can locally regularize the eigenvalues of the Hessian of the loss. In the context of least-squares linear regression, Behdin and Mazumder (2023) found that SAM exhibits lower bias and higher variance compared to gradient descent. However, all the above analyses of SAM utilize the Hessian information of the loss and require the smoothness property of the loss implicitly. The study for non-smooth neural networks, particularly for the classification task, remains open.

In this paper, our goal is to provide a theoretical basis demonstrating when SAM outperforms SGD. In particular, we consider a data distribution mainly characterized by the signal \(\mathbf{\mu}\) and input data dimension \(d\), and prove the following separation in terms of test error between SGD and SAM.

**Theorem 1.1** (Informal statement of Theorems 3.2 and 4.1).: _Let \(p\) be the strength of the label flipping noise. For any \(\epsilon>0\), under certain regularity conditions, with high probability, there exists \(0\leq t\leq T\) such that the training loss converges, i.e., \(L_{S}(\mathbf{W}^{(t)})\leq\epsilon\). Besides,_

1. _For SGD, when the signal strength_ \(\|\mathbf{\mu}\|_{2}\geq\Omega(d^{1/4})\)_, we have_ \(L_{\mathcal{D}}^{0-1}(\mathbf{W}^{(t)})\leq p+\epsilon\)_. When the signal strength_ \(\|\mathbf{\mu}\|_{2}\leq O(d^{1/4})\)_, we have_ \(L_{\mathcal{D}}^{0-1}(\mathbf{W}^{(t)})\geq p+0.1\)_._
2. _For SAM, provided the signal strength_ \(\|\mathbf{\mu}\|_{2}\geq\widehat{\Omega}(1)\)_, we have_ \(L_{\mathcal{D}}^{0-1}(\mathbf{W}^{(t)})\leq p+\epsilon\)_._

Our contributions are summarized as follows:

* We discuss how the loss landscape of two-layer convolutional ReLU networks is different from the smooth loss landscape and thus the current explanation for the success of SAM based on the Hessian information is insufficient for neural networks.
* To understand the limit of SGD, we precisely characterize the conditions under which benign overfitting can occur in training two-layer convolutional ReLU networks with SGD. To the best of our knowledge, this is the first benign overfitting result for neural network trained with mini-batch SGD. We also prove a phase transition phenomenon for SGD, which is illustrated in Figure 1.
* Under the conditions when SGD leads to harmful overfitting, we formally prove that SAM can achieve benign overfitting. Consequently, we establish a rigorous theoretical distinction between SAM and SGD, demonstrating that SAM strictly outperforms SGD in terms of generalization error. Specifically, we show that SAM effectively mitigates noise learning in the early stages of training, enabling neural networks to learn features more efficiently.

**Notation.** We use lower case letters, lower case bold face letters, and upper case bold face letters to denote scalars, vectors, and matrices respectively. For a vector \(\mathbf{v}=(v_{1},\cdots,v_{d})^{\top}\), we denote

Figure 1: Illustration of the phase transition between benign overfitting and harmful overfitting. The yellow region represents the regime under which the overfitted CNN trained by SGD is guaranteed to have a small excess risk, and the blue region represents the regime under which the excess risk is guaranteed to be a constant order (e.g., greater than \(0.1\)). The gray region is the regime where the excess risk is not characterized.

by \(\|\mathbf{v}\|_{2}:=\left(\sum_{j=1}^{d}v_{j}^{2}\right)^{1/2}\) its \(2\)-norm. For two sequence \(\{a_{k}\}\) and \(\{b_{k}\}\), we denote \(a_{k}=O(b_{k})\) if \(|a_{k}|\leq C|b_{k}|\) for some absolute constant \(C\), denote \(a_{k}=\Omega(b_{k})\) if \(b_{k}=O(a_{k})\), and denote \(a_{k}=\Theta(b_{k})\) if \(a_{k}=O(b_{k})\) and \(a_{k}=\Omega(b_{k})\). We also denote \(a_{k}=o(b_{k})\) if \(\lim|a_{k}/b_{k}|=0\). Finally, we use \(\widetilde{O}(\cdot)\) and \(\widetilde{\Omega}(\cdot)\) to omit logarithmic terms in the notation. We denote the set \(\{1,\cdots,n\}\) by \([n]\), and the set \(\{0,\cdots,n-1\}\) by \(\overline{[n]}\), respectively. The carnality of a set \(S\) is denoted by \(|S|\).

## 2 Preliminaries

### Data distribution

Our focus is on binary classification with label \(y\in\{\pm 1\}\). We consider the following data model, which can be seen as a special case of sparse coding model (Olshausen and Field, 1997; Allen-Zhu and Li, 2022; Ahn et al., 2022).

**Definition 2.1**.: Let \(\mathbf{\mu}\in\mathbb{R}^{d}\) be a fixed vector representing the signal contained in each data point. Each data point \((\mathbf{x},y)\) with input \(\mathbf{x}=[\mathbf{x}^{(1)\top},\mathbf{x}^{(2)\top},\ldots,\mathbf{x}^{(P) \top}]^{\top}\in\mathbb{R}^{P\times d},\mathbf{x}^{(1)},\mathbf{x}^{(2)}, \ldots,\mathbf{x}^{(P)}\in\mathbb{R}^{d}\) and label \(y\in\{-1,1\}\) is generated from a distribution \(\mathcal{D}\) specified as follows:

1. The true label \(\widehat{y}\) is generated as a Rademacher random variable, i.e., \(\mathbb{P}[\widehat{y}=1]=\mathbb{P}[\widehat{y}=-1]=1/2\). The observed label \(y\) is then generated by flipping \(\widehat{y}\) with probability \(p\) where \(p<1/2\), i.e., \(\mathbb{P}[y=\widehat{y}]=1-p\) and \(\mathbb{P}[y=-\widehat{y}]=p\).
2. A noise vector \(\mathbf{\xi}\) is generated from the Gaussian distribution \(\mathcal{N}(\mathbf{0},\sigma_{p}^{2}\mathbf{I})\), where \(\sigma_{p}^{2}\) is the variance.
3. One of \(\mathbf{x}^{(1)},\mathbf{x}^{(2)},\ldots,\mathbf{x}^{(P)}\) is randomly selected and then assigned as \(y\cdot\mathbf{\mu}\), which represents the signal, while the others are given by \(\mathbf{\xi}\), which represents noises.

The data distribution in Definition 2.1 has also been extensively employed in several previous works (Allen-Zhu and Li, 2020; Jelassi and Li, 2022; Shen et al., 2022; Cao et al., 2022; Kou et al., 2023). When \(P=2\), this data distribution aligns with the one analyzed in Kou et al. (2023). This distribution is inspired by image data, where the input is composed of different patches, with only a few patches being relevant to the label. The model has two key vectors: the feature vector and the noise vector. For any input vector \(\mathbf{x}=[\mathbf{x}^{(1)},\ldots,\mathbf{x}^{(P)}]\), there is exactly one \(\mathbf{x}^{(j)}=y\mathbf{\mu}\), and the others are random Gaussian vectors. For example, the input vector \(\mathbf{x}\) can be \([y\mathbf{\mu},\mathbf{\xi},\ldots,\mathbf{\xi}],\ldots,[\mathbf{\xi},\ldots,y\mathbf{\mu},\mathbf{\xi}]\) or \([\mathbf{\xi},\ldots,\mathbf{\xi},y\mathbf{\mu}]\): the signal patch \(y\mathbf{\mu}\) can appear at any position. To avoid harmful overfitting, the model must learn the feature vector rather than the noise vector.

### Neural Network and Training Loss

To effectively learn the distribution as per Definition 2.1, it is advantageous to utilize a shared weights structure, given that the specific signal patch is not known beforehand. When \(P>n\), shared weights become indispensable as the location of the signal patch in the test can differ from the location of the signal patch in the training data.

We consider a two-layer convolutional neural network whose filters are applied to the \(P\) patches \(\mathbf{x}_{1},\cdots,\mathbf{x}_{P}\) separately, and the second layer parameters of the network are fixed as \(+1/m\) and \(-1/m\) respectively, where \(m\) is the number of convolutional filters. Then the network can be written as \(f(\mathbf{W},\mathbf{x})=F_{+1}(\mathbf{W}_{+1},\mathbf{x})-F_{-1}(\mathbf{W}_ {-1},\mathbf{x})\), where \(F_{+1}(\mathbf{W}_{+1},\mathbf{x})\) and \(F_{-1}(\mathbf{W}_{-1},\mathbf{x})\) are defined as

\[F_{j}(\mathbf{W}_{j},\mathbf{x})=m^{-1}{\sum_{r=1}^{m}{\sum_{p=1}^{P}\sigma( \langle\mathbf{w}_{j,r},\mathbf{x}^{(p)}\rangle)}}. \tag{1}\]

Here we consider ReLU activation function \(\sigma(z)=\mathds{1}(z\geq 0)z\), \(\mathbf{w}_{j,r}\in\mathbb{R}^{d}\) denotes the weight for the \(r\)-th filter, and \(\mathbf{W}_{j}\) is the collection of model weights associated with \(F_{j}\) for \(j=\pm 1\). We use \(\mathbf{W}\) to denote the collection of all model weights. Denote the training data set by \(\mathcal{S}=\{(\mathbf{x}_{i},y_{i})\}_{i\in[n]}\). We train the above CNN model by minimizing the empirical cross-entropy loss function

\[L_{\mathcal{S}}(\mathbf{W})=n^{-1}{\sum_{i\in[n]}}\ell(y_{i}f(\mathbf{W}, \mathbf{x}_{i})),\]

where \(\ell(z)=\log(1+\exp(-z))\).

### Training Algorithm

**Minibatch Stochastic Gradient Descent.** For epoch \(t\), the training data set \(S\) is randomly divided into \(H:=n/B\) mini batches \(\mathcal{I}_{t,b}\) with batch size \(B\geq 2\). The empirical loss for batch \(\mathcal{I}_{t,b}\) is defined as \(L_{\mathcal{I}_{t,b}}(\mathbf{W})=(1/B)\sum_{i\in\mathcal{I}_{t,b}}\ell(y_{i}f( \mathbf{W},\mathbf{x}_{i}))\). Then the gradient descent update of the filters in the CNN can be written as

\[\mathbf{w}^{(t,b+1)}=\mathbf{w}^{(t,b)}-\eta\cdot\nabla_{\mathbf{W}}L_{ \mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)}), \tag{2}\]

where the gradient of the empirical loss \(\nabla_{\mathbf{W}}L_{\mathcal{I}_{t,b}}\) is the collection of \(\nabla_{\mathbf{w}_{j},r}L_{\mathcal{I}_{t,b}}\) as follows

\[\nabla_{\mathbf{w}_{j},r}L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)}) =\frac{(P-1)}{Bm}\sum_{i\in\mathcal{I}_{t,b}}\ell_{i}^{\prime(t,b )}\cdot\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle) \cdot jy_{i}\mathbf{\xi}_{i}\] \[\qquad+\frac{1}{Bm}\sum_{i\in\mathcal{I}_{t,b}}\ell_{i}^{\prime(t,b)}\cdot\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\widehat{y}_{i}\mathbf{ \mu}\rangle)\cdot\widehat{y}_{i}y_{i}j\mathbf{\mu}, \tag{3}\]

for all \(j\in\{\pm 1\}\) and \(r\in[m]\). Here we introduce a shorthand notation \(\ell_{i}^{\prime(t,b)}=\ell^{\prime}[y_{i}\cdot f(\mathbf{W}^{(t,b)},\mathbf{x }_{i})]\) and assume the gradient of the ReLU activation function at \(0\) to be \(\sigma^{\prime}(0)=1\) without loss of generality. We use \((t,b)\) to denote epoch index \(t\) with mini-batch index \(b\) and use \((t)\) as the shorthand of \((t,0)\). We initialize SGD by random Gaussian, where all entries of \(\mathbf{W}^{(0)}\) are sampled from i.i.d. Gaussian distributions \(\mathcal{N}(0,\sigma_{0}^{2})\), with \(\sigma_{0}^{2}\) being the variance. From (3), we can infer that the loss landscape of the empirical loss is highly non-smooth because the ReLU function is not differentiable at zero. In particular, when \(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}\rangle\) is close to zero, even a very small perturbation can greatly change the activation pattern \(\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}\rangle)\) and thus change the direction of \(\nabla_{\mathbf{w}_{j},r}L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})\). This observation prevents the analysis technique based on the Taylor expansion with the Hessian matrix, and calls for a more sophisticated activation pattern analysis.

**Sharpness Aware Minimization.** Given an empirical loss function \(L_{S}(\mathbf{W})\) with trainable parameter \(\mathbf{W}\), the idea of SAM is to minimize a perturbed empirical loss at the worst point in the neighborhood ball of \(\mathbf{W}\) to ensure a uniformly low training loss value. In particular, it aims to solve the following optimization problem

\[\min_{\mathbf{W}}L_{S}^{\text{SAM}}(\mathbf{W}),\quad\text{where} \quad L_{S}^{\text{SAM}}(\mathbf{W}):=\max_{\|\mathbf{\epsilon}\|_{2}\leq\tau}L_{S }(\mathbf{W}+\mathbf{\epsilon}), \tag{4}\]

where the hyperparameter \(\tau\) is called the perturbation radius. However, directly optimizing \(L_{S}^{\text{SAM}}(\mathbf{W})\) is computationally expensive. In practice, people use the following sharpness-aware minimization (SAM) algorithm (Foret et al., 2020; Zheng et al., 2021) to minimize \(L_{S}^{\text{SAM}}(\mathbf{W})\) efficiently,

\[\mathbf{W}^{(t+1)}=\mathbf{W}^{(t)}-\eta\nabla_{\mathbf{W}}L_{S}\big{(} \mathbf{W}+\widehat{\mathbf{\epsilon}}\big{)},\quad\text{where}\quad\widehat{ \mathbf{\epsilon}}=\tau\cdot\frac{\nabla_{\mathbf{W}}L_{S}(\mathbf{W})}{\|\nabla_ {\mathbf{W}}L_{S}(\mathbf{W})\|_{F}}. \tag{5}\]

When applied to SGD in (2), the gradient \(\nabla_{\mathbf{W}}L_{S}\) in (5) is further replaced by stochastic gradient \(\nabla_{\mathbf{W}}L_{\mathcal{I}_{t,b}}\)(Foret et al., 2020). The detailed algorithm description of SAM in shown in Algorithm 1.

```
Input: Training set \(\mathcal{S}=\cup_{i=1}^{n}\{(\mathbf{x}_{i},\mathbf{y}_{i})\}\), Batch size \(B\), step size \(\eta>0\), neighborhood size \(\tau>0\).  Initialize weights \(\mathbf{W}^{(0)}\). for\(t=0,1,\dots,T-1\)do  Randomly divide the training data set into \(H\) mini batches \(\{\mathcal{I}_{t,b}\}_{b=0}^{H-1}\). for\(b=0,1,\dots,H-1\)do  We calculate the perturbation \(\widehat{\mathbf{\epsilon}}^{(t,b)}=\tau\frac{\nabla_{\mathbf{W}}L_{\mathcal{I}_{t,b }}(\mathbf{W}^{(t,b)})}{\|\nabla_{\mathbf{W}}L_{\mathcal{I}_{t,b}}(\mathbf{W} ^{(t,b)})\|_{F}}\).  Update model parameters: \(\mathbf{W}^{(t,b+1)}=\mathbf{W}^{(t,b)}-\eta\nabla_{\mathbf{W}}L_{\mathcal{I} _{t,b}}(\mathbf{W})|_{\mathbf{W}=\mathbf{W}^{(t,b)}+\widehat{\mathbf{\epsilon}}^{(t,b)}}\). endfor  Update model parameters: \(\mathbf{W}^{(t+1,0)}=\mathbf{W}^{(t,H)}\) endfor
```

**Algorithm 1** Minibatch Sharpness Aware Minimization

## 3 Result for SGD

In this section, we present our main theoretical results for the CNN trained with SGD. Our results are based on the following conditions on the dimension \(d\), sample size \(n\), neural network width \(m\), initialization scale \(\sigma_{0}\) and learning rate \(\eta\).

**Condition 3.1**.: Suppose there exists a sufficiently large constant \(C\), such that the following hold:

1. Dimension \(d\) is sufficiently large: \(d\geq\widetilde{\Omega}\Big{(}\max\{nP^{-2}\sigma_{p}^{-2}\|\mathbf{\mu}\|_{2}^{2},n^{ 2},P^{-2}\sigma_{p}^{-2}Bm\}\Big{)}\).
2. Training sample size \(n\) and neural network width satisfy \(m,n\geq\widetilde{\Omega}(1)\).
3. The \(2\)-norm of the signal satisfies \(\|\mathbf{\mu}\|_{2}\geq\widetilde{\Omega}(P\sigma_{p})\).
4. The noise rate \(p\) satisfies \(p\leq 1/C\).
5. The standard deviation of Gaussian initialization \(\sigma_{0}\) is appropriately chosen such that \(\sigma_{0}\leq\widetilde{O}\Big{(}\big{(}\max\big{\{}P\sigma_{p}d/\sqrt{n},\| \mathbf{\mu}\|_{2}\big{\}}\big{)}^{-1}\Big{)}\).
6. The learning rate \(\eta\) satisfies \(\eta\leq\widetilde{O}\Big{(}\big{(}\max\big{\{}P^{2}\sigma_{p}^{2}d^{3/2}/(Bm ),P^{2}\sigma_{p}^{2}d/B,n\|\mu\|_{2}/(\sigma_{0}B\sqrt{d}m),\)

\(nP\sigma_{p}\|\mathbf{\mu}\|_{2}/(B^{2}m\epsilon)\big{\}}\Big{)}^{-1}\Big{)}\).

The conditions imposed on the data dimensions \(d\), network width \(m\), and the number of samples \(n\) ensure adequate overparameterization of the network. Additionally, the condition on the learning rate \(\eta\) facilitates efficient learning by our model. By concentration inequality, with high probability, the \(\ell_{2}\) norm of the noise patch is of order \(\Theta(d\sigma_{p}^{2})\). Therefore, the quantity \(d\sigma_{p}^{2}\) can be viewed as the strength of the noise. Comparable conditions have been established in Chatterji and Long (2021); Cao et al. (2022); Frei et al. (2022); Kou et al. (2023). Based on the above condition, we first present a set of results on benign/harmful overfitting for SGD in the following theorem.

**Theorem 3.2** (Benign/harmful overfitting of SGD in training CNNs).: _For any \(\epsilon>0\), under Condition 3.1, with probability at least \(1-\delta\) there exists \(t=\widetilde{O}(\eta^{-1}\epsilon^{-1}mnd^{-1}P^{-2}\sigma_{p}^{-2})\) such that:_

1. _The training loss converges, i.e.,_ \(L_{S}(\mathbf{W}^{(t)})\leq\epsilon\)_._
2. _When_ \(n\|\mathbf{\mu}\|_{2}^{4}\geq C_{1}dP^{4}\sigma_{p}^{4}\)_, the test error_ \(L_{\mathcal{D}}^{0-1}(\mathbf{W}^{(t)})\leq p+\epsilon\)_._
3. _When_ \(n\|\mathbf{\mu}\|_{2}^{4}\leq C_{3}dP^{4}\sigma_{p}^{4}\)_, the test error_ \(L_{\mathcal{D}}^{0-1}(\mathbf{W}^{(t)})\geq p+0.1\)_._

Theorem 3.2 reveals a sharp phase transition between benign and harmful overfitting for CNN trained with SGD. This transition is determined by the relative scale of the signal strength and the data dimension. Specifically, if the signal is relatively large such that \(n\|\mathbf{\mu}\|_{2}^{4}\geq C_{1}d(P-1)^{4}\sigma_{p}^{4}\), the model can efficiently learn the signal. As a result, the test error decreases, approaching the Bayes risk \(p\), although the presence of label flipping noise prevents the test error from reaching zero. Conversely, when the condition \(n\|\mathbf{\mu}\|_{2}^{4}\leq C_{3}d(P-1)^{4}\sigma_{p}^{4}\) holds, the test error fails to approach the Bayes risk. This phase transition is empirically illustrated in Figure 2. In both scenarios, the model is capable of fitting the training data thoroughly, even for examples with flipped labels. This finding aligns with longstanding empirical observations.

The negative result of SGD, which encompasses the third point of Theorem 3.2 and the high test error observed in Figure 2, suggests that the signal strength needs to scale with the data dimension to enable benign overfitting. This constraint substantially undermines the efficiency of SGD, particularly when dealing with high-dimensional data. A significant part of this limitation stems from the fact that SGD does not inhibit the model from learning noise, leading to a comparable rate of signal and noise learning during iterative model parameter updates. This inherent limitation of SGD is effectively addressed by SAM, as we will discuss later in Section 4.

### Analysis of Mini-Batch SGD

In contrast to GD, SGD does not utilize all the training data at each iteration. Consequently, different samples may contribute to parameters differently, leading to possible unbalancing in parameters. To analyze SGD, we extend the signal-noise decomposition technique developed by Kou et al. (2023); Cao et al. (2022) for GD, which in our case is formally defined as:

**Lemma 3.3**.: _Let \(\mathbf{w}_{j,r}^{(t,b)}\) for \(j\in\{\pm 1\}\), \(r\in[m]\) be the convolution filters of the CNN at the \(b\)-th batch of \(t\)-th epoch of gradient descent. Then there exist unique coefficients \(\gamma_{j,r}^{(t,b)}\) and \(p_{j,r,i}^{(t,b)}\) such that_

\[\mathbf{w}_{j,r}^{(t,b)}=\mathbf{w}_{j,r}^{(0,0)}+j\cdot\gamma_{j,r}^{(t,b)} \cdot\|\mathbf{\mu}\|_{2}^{-2}\cdot\mathbf{\mu}+\frac{1}{P-1}\sum_{i=1}^{n}\rho_{j,r,i} ^{(t,b)}\cdot\|\mathbf{\xi}_{i}\|_{2}^{-2}\cdot\mathbf{\xi}_{i}. \tag{6}\]

_Further denote \(\overline{\rho}_{j,r,i}^{(t,b)}:=\rho_{j,r,i}^{(t,b)}\mathds{1}(\rho_{j,r,i}^{( t,b)}\geq 0)\), \(\rho_{j,r,i}^{(t,b)}:=\rho_{j,r,i}^{(t,b)}\mathds{1}(\rho_{j,r,i}^{(t,b)}\leq 0)\). Then_

\[\mathbf{w}_{j,r}^{(t,b)}=\mathbf{w}_{j,r}^{(0,0)}+j\gamma_{j,r}^{(t,b)}\|\mathbf{ \mu}\|_{2}^{-2}\mathbf{\mu}+\frac{1}{P-1}\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(t,b)}\|\mathbf{\xi}_{i}\|_{2}^{-2}\mathbf{\xi}_{i}+\frac{1}{P-1}\sum_{i=1}^{n}\underline {\rho}_{j,r,i}^{(t,b)}\|\mathbf{\xi}_{i}\|_{2}^{-2}\mathbf{\xi}_{i}. \tag{7}\]

Note that (7) is a variant of (6): by decomposing the coefficient \(\rho_{j,r,i}^{(t,b)}\) into \(\overline{\rho}_{j,r,i}^{(t,b)}\) and \(\underline{\rho}_{j,r,i}^{(t,b)}\), we can streamline our proof process. The normalization terms \(\frac{1}{P-1}\), \(\|\mathbf{\mu}\|_{2}^{-2}\), and \(\|\mathbf{\xi}_{i}\|_{2}^{-2}\) ensure that \(\gamma_{j,r}^{(t,b)}\approx\langle\mathbf{w}_{j,r}^{(t,b)},\mu\rangle\) and \(\rho_{j,r}^{(t,b)}\approx(P-1)\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle\). Through signal-noise decomposition, we characterize the learning progress of signal \(\mathbf{\mu}\) using \(\gamma_{j,r}^{(t,b)}\), and the learning progress of noise using \(\rho_{j,r}^{(t,b)}\). This decomposition turns the analysis of SGD updates into the analysis of signal noise coefficients. Kou et al. (2023) extend this technique to the ReLU activation function as well as in the presence of label flipping noise. However, mini-batch SGD updates amplify the complications introduced by label flipping noise, making it more difficult to ensure learning. We have developed advanced methods for coefficient balancing and activation pattern analysis. These techniques will be thoroughly discussed in the sequel. The progress of signal learning is characterized by \(\gamma_{j,r}^{(t,b)}\), whose update rule is as follows:

\[\gamma_{j,r}^{(t,b+1)}=\gamma_{j,r}^{(t,b)}-\frac{\eta}{Bm}\cdot \left[\sum_{i\in\mathcal{I}_{t,b}\cap\mathcal{S}_{+}}\ell_{i}^{\prime(t,b)} \sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\widehat{y}_{i}\cdot\mathbf{\mu} \rangle)\right.\] \[\left.-\sum_{i\in\mathcal{I}_{t,b}\cap\mathcal{S}_{-}}\ell_{i}^{ \prime(t,b)}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\widehat{y}_{i} \cdot\mathbf{\mu}\rangle)\right]\cdot\|\mathbf{\mu}\|_{2}^{2}. \tag{8}\]

Here, \(\mathcal{I}_{t,b}\) represents the indices of samples in batch \(b\) of epoch \(t\), \(S_{+}\) denotes the set of clean samples where \(y_{i}=\widehat{y}_{i}\), and \(S_{-}\) represents the set of noisy samples where \(y_{i}=-\widehat{y}_{i}\). The updates of \(\gamma_{j,r}^{(t,b)}\) comprise an increment arising from sample learning, counterbalanced by a decrement due to noisy sample learning. Both empirical and theoretical analyses have demonstrated that overparametrization allows the model to fit even random labels. This occurs when the negative term \(\sum_{i\in\mathcal{I}_{t,b}\cap\mathcal{S}_{-}}\ell_{i}^{\prime(t,b)}\sigma^{ \prime}(\langle\mathbf{w}_{j,r}^{(t,b)},y_{i}\cdot\mathbf{\mu}\rangle)\) primarily drives model learning. Such unfavorable scenarios can be attributed to two possible factors. Firstly, the gradient of the loss \(\ell_{i}^{\prime(t,b)}\) might be significantly larger for noisy samples compared to clean samples. Secondly, during certain epochs, the majority of samples may be noisy, meaning that \(\mathcal{I}_{t,b}\cap S_{-}\) significantly outnumbers \(\mathcal{I}_{t,b}\cap\mathcal{S}_{+}\).

To deal with the first factor, we have to control the ratio of the loss gradient with regard to different samples, as depicted in (9). Given that noisy samples may overwhelm a single batch, we impose an additional requirement: the ratio of the loss gradient must be controllable across different batches within a single epoch, i.e.,

\[\ell_{i}^{\prime(t,b_{1})}/\ell_{k}^{\prime(t,b_{2})}\leq C_{2}. \tag{9}\]

As \(\ell^{\prime}(z_{1})/\ell^{\prime}(z_{2})\approx\exp(z_{2}-z_{1})\), we can upper bound \(\ell_{i}^{\prime(t,b_{1})}/\ell_{k}^{\prime(t,b_{2})}\) by \(y_{i}\cdot f(\mathbf{W}^{(t,b_{1})},\mathbf{x}_{i})-y_{k}\cdot f(\mathbf{W}^{(t,b_{2})},\mathbf{x}_{k})\), which can be further upper bounded by \(\sum_{r}\overline{\rho}_{y_{i},r,i}^{(t,b_{1})}-\sum_{r}\overline{\rho}_{y_{i},r,k}^{(t,b_{2})}\) with a small error. Therefore, the proof of (9) can be reduced to proving a uniform bound on the difference among \(\overline{\rho}_{y_{i},r,i}^{(t,b_{1})}\), i.e., \(\sum_{r=1}^{m}\overline{\rho}_{y_{i},r,i}^{(t,b_{1})}-\sum_{r=1}^{m}\overline{ \rho}_{y_{k},r,k}^{(t,b_{2})}\leq\kappa,\;\forall\;i,k\).

However, achieving this uniform upper bound turns out to be challenging, since the updates of \(\overline{\rho}_{j,r,i}^{(t,b)}\)'s are not evenly distributed across different batches within an epoch. Each mini-batch update utilizes only a portion of the samples, meaning that some \(\overline{\rho}_{y_{i},r,i}^{(t,b)}\) can increase or decrease much more thanthe others. Therefore, the uniformly bounded difference can only be achieved after the entire epoch is processed. Consequently, we have to first bound the difference among \(\overline{p}^{(t,b)}_{y_{i},r,i}\)'s after each entire epoch, and then control the maximal difference within one epoch. The full batch (epoch) update rule is established as follows:

\[\sum_{r=1}^{m}\big{[}\overline{p}^{(t+1,0)}_{y_{i},r,i}-\overline{ p}^{(t+1,0)}_{y_{k},r,k}\big{]} =\sum_{r=1}^{m}\big{[}\overline{p}^{(t,0)}_{y_{i},r,i}-\overline{ p}^{(t,0)}_{y_{k},r,k}\big{]}-\frac{\eta(P-1)^{2}}{Bm}\cdot\Big{(}|\widetilde{S }^{(t,b^{(t)}_{i}}_{i}|\ell^{\prime(t,b^{(t)}_{i}}_{i})\cdot\|\mathbf{\xi}_{i}\|_{2 }^{2}\] \[\qquad\qquad\qquad\qquad-|\widetilde{S}^{(t,b^{(t)}_{k}}_{k})| \ell^{\prime(t,b^{(t)}_{k}}_{k})\cdot\|\mathbf{\xi}_{k}\|_{2}^{2}\Big{)}. \tag{10}\]

Here, \(b^{(t)}_{i}\) denotes the batch to which sample \(i\) belongs in epoch \(t\), and \(\widetilde{S}^{(t,b^{(t)}_{i}}_{i})\) represents the parameters that learn \(\mathbf{\xi}_{i}\) at epoch \(t\) defined as

\[\widetilde{S}^{(t,b)}_{i}:=\{r:\langle\mathbf{w}^{(t,b)}_{y_{i},r },\mathbf{\xi}_{i}\rangle>0\}. \tag{11}\]

Therefore, the update of \(\sum_{r=1}^{m}\big{[}\overline{p}^{(t,0)}_{y_{i},r,i}-\overline{p}^{(t,0)}_{y_ {k},r,k}\big{]}\) is indeed characterized by the activation pattern of parameters, which serves as the key technique for analyzing the full batch update of \(\sum_{r=1}^{m}\big{[}\overline{\rho}^{(t,0)}_{y_{i},r,i}-\overline{\rho}^{(t,0 )}_{y_{k},r,k}\big{]}\). However, analyzing the pattern of \(S^{(t,b)}_{i}\) directly is challenging since \(\langle\mathbf{w}^{(t,b)}_{y_{i},r},\mathbf{\xi}_{i}\rangle\) fluctuates in mini-batches without sample \(i\). Therefore, we introduce the set series \(S^{(t,b)}_{i}\) as the activation pattern with certain threshold as follows:

\[S^{(t,b)}_{i}:=\{r:\langle\mathbf{w}^{(t,b)}_{y_{i},r},\mathbf{\xi}_ {i}\rangle>\sigma_{0}\sigma_{p}\sqrt{d}/\sqrt{2}\}. \tag{12}\]

The following lemma suggests that the set of activated parameters \(S^{(t,0)}_{i}\) is a non-decreasing sequence with regards to \(t\), and the set of plain activated parameters \(\widetilde{S}^{(t,b)}_{i}\) always include \(S^{(t,0)}_{i}\). Consequently, \(S^{(0,0)}_{i}\) is always included in \(\widetilde{S}^{(t,b)}_{i}\), guaranteeing that \(\mathbf{\xi}_{i}\) can always be learned by some parameter. And this further makes sure the difference among \(\overline{\rho}^{(t,b)}_{y_{i},r,i}\) is bounded, as well as \(\ell^{\prime(t,b_{1})}_{i}/\ell^{\prime(t,b_{2})}_{k}\leq C_{2}\). In the proof for SGD, we consider the learning period \(0\leq t\leq T^{*}\), where \(T^{*}=\eta^{-1}\mathrm{poly}(\epsilon^{-1},d,n,m)\) is the maximum number of admissible iterations.

**Lemma 3.4** (Informal Statement of Lemma C.8).: _For all \(t\in[0,T^{*}]\) and \(b\in\overline{[H]}\), we have_

\[S^{(t-1,0)}_{i}\subseteq S^{(t,0)}_{i}\subseteq\widetilde{S}^{ (t,b)}_{i}. \tag{13}\]

As we have mentioned above, if noisy samples outnumber clean samples, \(\gamma^{(t,b)}_{j,r}\) may also decrease. To deal with such a scenario, we establish a two-stage analysis of \(\gamma^{(t,b)}_{j,r}\) progress. In the first stage, when \(-\ell^{\prime}_{i}\) is lower bound by a positive constant, we prove that there are enough batches containing sufficient clear samples. This is characterized by the following high-probability event.

**Lemma 3.5** (Informal Statement of Lemma B.6).: _With high probability, for all \(T\in[\widetilde{O}(1),T^{*}]\), there exist at least \(c_{1}\cdot T\) epochs among \([0,T]\), such that at least \(c_{2}\cdot H\) batches in each of these epochs satisfying the following condition:_

\[|S_{+}\cap S_{y}\cap\mathcal{I}_{t,b}|\in[0.25B,0.75B]. \tag{14}\]

After the first stage of \(T=\Theta(\eta^{-1}m(P-1)^{-2}\sigma_{p}^{-2}d^{-1})\) epochs, we have \(\gamma^{(T,0)}_{j,r}=\Omega\Big{(}n\frac{\|\mathbf{\mu}\|_{2}^{2}}{(P-1)^{2} \sigma_{p}^{2}d}\Big{)}\). The scale of \(\gamma^{(T,0)}_{j,r}\) guarantees that \(\langle\mathbf{w}^{(t,b)}_{j,r},\mathbf{\mu}\rangle\) remains resistant to intra-epoch fluctuations. Consequently, this implies the sign of \(\langle\mathbf{w}^{(t,b)}_{j,r},\mathbf{\mu}\rangle\) will persist unchanged throughout the entire epoch. Without loss of generality, we suppose that \(\langle\mathbf{w}^{(t,b)}_{j,r},\mathbf{\mu}\rangle>0\), then the update of \(\gamma^{(t,b)}_{j,r}\) can be written as follows:

\[\gamma^{(t+1,0)}_{j,r}=\gamma^{(t,0)}_{j,r}+\frac{\eta}{Bm}\cdot\bigg{[}\min _{i\in\mathcal{I}_{t,b},b}|\ell^{\prime(t,b)}_{i}||S_{+}\cap S_{1}|-\max_{i \in\mathcal{I}_{t,b},b}|\ell^{\prime(t,b)}_{i}||S_{-}\cap S_{-1}|\bigg{]} \cdot\|\mathbf{\mu}\|_{2}^{2}. \tag{15}\]

As we have proved the balancing of logits \(\ell^{\prime(t,b)}_{i}\) across batches, the progress analysis of \(\gamma^{(t+1,0)}_{j,r}\) is established to characterize the signal learning of SGD.

## 4 Result for SAM

In this section, we present the positive results for SAM in the following theorem.

**Theorem 4.1**.: _For any \(\epsilon>0\), under Condition 3.1 with \(\sigma_{0}=\widetilde{\Theta}(P^{-1}\sigma_{p}^{-1}d^{-1/2})\), choose \(\tau=\Theta\Big{(}\frac{m\sqrt{B}}{P\sigma_{p}\sqrt{d}}\Big{)}\). With probability at least \(1-\delta\), neural networks first trained with SAM with \(O\Big{(}\eta^{-1}\epsilon^{-1}n^{-1}mB\|\mathbf{\mu}\|_{2}^{-2}\Big{)}\) iterations, then trained with SGD with \(\widetilde{O}\Big{(}\eta^{-1}\epsilon^{-1}mnd^{-1}P^{-2}\sigma_{p}^{-2}\Big{)}\) iterations can find \(\mathbf{W}^{(t)}\) such that,_

1. _The training loss satisfies_ \(L_{S}(\mathbf{W}^{(t)})\leq\epsilon\)_._
2. _The test error_ \(L_{\mathcal{D}}^{0-1}(\mathbf{W}^{(t)})\leq p+\epsilon\)_._

In contrast to Theorem 3.2, Theorem 4.1 demonstrates that CNNs trained by SAM exhibit benign overfitting under much milder conditions. This condition is almost dimension-free, as opposed to the threshold of \(\|\mathbf{\mu}\|_{2}^{4}\geq\widetilde{\Omega}((d/n)P^{4}\sigma_{p}^{4})\) for CNNs trained by SGD. The discrepancy in the thresholds can be observed in Figure 1. This difference is because SAM introduces a perturbation during the model parameter update process, which effectively prevents the early-stage memorization of noise by deactivating the corresponding neurons.

### Noise Memorization Prevention

In this subsection, we will show how SAM can prevent noise memorization by changing the activation pattern of the neurons. For SAM, we have the following update rule of decomposition coefficients \(\gamma_{j,r}^{(t,b)},\overline{\rho}_{j,r,i}^{(t,b)},\underline{\rho}_{j,r,i}^ {(t,b)}\).

**Lemma 4.2**.: _The coefficients \(\gamma_{j,r}^{(t,b)},\overline{\rho}_{j,r,i}^{(t,b)},\underline{\rho}_{j,r,i}^ {(t,b)}\) defined in Lemma 3.3 satisfy the following iterative equations for all \(r\in[m]\), \(j\in\{\pm 1\}\) and \(i\in[n]\):_

\[\gamma_{j,r}^{(0,0)},\overline{\rho}_{j,r,i}^{(0,0)},\underline{ \rho}_{j,r,i}^{(0,0)}=0,\] \[\gamma_{j,r}^{(t,b+1)}=\gamma_{j,r}^{(t,b)}-\frac{\eta}{Bm} \cdot\bigg{[}\sum_{i\in\mathcal{I}_{t,b}\cap S_{+}}\ell_{i}^{(t,b)} \sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)}+\widehat{\mathbf{\epsilon}}_{j,r}^ {(t,b)},\widehat{y}_{i}\cdot\mathbf{\mu}\rangle)\] \[-\sum_{i\in\mathcal{I}_{t,b}\cap S_{-}}\ell_{i}^{(t,b)}\sigma^{ \prime}(\langle\mathbf{w}_{j,r}^{(t,b)}+\widehat{\mathbf{\epsilon}}_{j,r}^{(t,b)},\widehat{y}_{i}\cdot\mathbf{\mu}\rangle)\bigg{]}\cdot\|\mathbf{\mu}\|_{2}^{2},\] \[\overline{\rho}_{j,r,i}^{(t,b+1)}=\overline{\rho}_{j,r,i}^{(t,b)} -\frac{\eta(P-1)^{2}}{Bm}\cdot\ell_{i}^{\prime(t,b)}\cdot\sigma^{\prime}( \langle\mathbf{w}_{j,r}^{(t,b)}+\widehat{\mathbf{\epsilon}}_{j,r}^{(t,b)},\mathbf{ \xi}_{i}\rangle)\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}\cdot\mathds{1}(y_{i}=j)\, \mathds{1}(i\in\mathcal{I}_{t,b}),\] \[\underline{\rho}_{j,r,i}^{(t,b+1)}=\underline{\rho}_{j,r,i}^{(t,b )}+\frac{\eta(P-1)^{2}}{Bm}\cdot\ell_{i}^{\prime(t,b)}\cdot\sigma^{\prime}( \langle\mathbf{w}_{j,r}^{(t,b)}+\widehat{\mathbf{\epsilon}}_{j,r}^{(t,b)},\mathbf{\xi} _{i}\rangle)\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}\cdot\mathds{1}(y_{i}=-j)\,\mathds{1 }(i\in\mathcal{I}_{t,b}),\]

_where \(\mathcal{I}_{t,b}\) denotes the sample index set of the \(b\)-th batch in the \(t\)-th epoch._

The primary distinction between SGD and SAM is how neuron activation is determined. In SAM, the activation is based on the perturbed weight \(\mathbf{w}_{j,r}^{(t,b)}+\widehat{\mathbf{\epsilon}}_{j,r}^{(t,b)}\), whereas in SGD, it is determined by the unperturbed weight \(\mathbf{w}_{j,r}^{(t,b)}\). This perturbation to the weight update process at each iteration gives SAM an intriguing denoising property. Specifically, if a neuron is activated by noise in the SGD update, it will subsequently become deactivated after the perturbation, as stated in the following lemma.

**Lemma 4.3** (Informal Statement of Lemma d.5).: _Suppose the Condition 3.1 holds with parameter choices in Theorem 4.1, if \(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{k}\rangle\geq 0\), \(k\in\mathcal{I}_{t,b}\) and \(j=y_{k}\), then \(\langle\mathbf{w}_{j,r}^{(t,b)}+\widehat{\mathbf{\epsilon}}_{j,r}^{(t,b)},\mathbf{\xi}_ {k}\rangle<0\)._

By leveraging this intriguing property, we can derive a constant upper bound for the noise coefficients \(\overline{\rho}_{j,r,i}^{(t,b)}\) by considering the following cases:

1. If \(\mathbf{\xi}_{i}\) is not in the current batch, then \(\overline{\rho}_{j,r,i}^{(t,b)}\) will not be updated in the current iteration.

2. If \(\mathbf{\xi}_{i}\) is in the current batch, we discuss two cases: 1. If \(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle\geq 0\), then by Lemma 4.3, one can know that \(\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)}+\tilde{\mathbf{\varepsilon}}_{j,r}^ {(t,b)},\mathbf{\xi}_{i}\rangle)=0\) and thus \(\overline{\rho}_{j,r,i}^{(t,b)}\) will not be updated in the current iteration. 2. If \(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle\leq 0\), then given that \(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle\approx\overline{\rho}_{j, r,i}^{(t,b)}\) and \(\overline{\rho}_{j,r,i}^{(t,b+1)}\leq\overline{\rho}_{j,r,i}^{(t,b)}+\frac{\eta( P-1)^{2}\|\mathbf{\xi}_{i}\|_{2}^{2}}{Bm}\), we can assert that, provided \(\eta\) is sufficiently small, the term \(\overline{\rho}_{j,r,i}^{(t,b)}\) can be upper bounded by a small constant.

In contrast to the analysis of SGD, which provides an upper bound for \(\overline{\rho}_{j,r,i}^{(t,b)}\) of order \(O(\log d)\), the noise memorization prevention property described in Lemma 4.3 allows us to obtain an upper bound for \(\overline{\rho}_{j,r,i}^{(t,b)}\) of order \(O(1)\) throughout \([0,T_{1}]\). This indicates that SAM memorizes less noise compared to SGD. On the other hand, the signal coefficient \(\gamma_{j,r,i}^{(t)}\) also increases to \(\Omega(1)\) for SAM, following the same argument as in SGD. This property ensures that training with SAM does not exhibit harmful overfitting for the same signal-to-noise ratio at which training with SGD suffers from harmful overfitting.

## 5 Experiments

In this section, we conduct synthetic experiments to validate our theory. Additional experiments on real data sets can be found in Appendix A.

We set training data size \(n=20\) without label-flipping noise. Since the learning problem is rotation-invariant, without loss of generality, we set \(\mathbf{\mu}=\|\mathbf{\mu}\|_{2}\cdot[1,0,\ldots,0]^{\top}\). We then generate the noise vector \(\mathbf{\xi}\) from the Gaussian distribution \(\mathcal{N}(\mathbf{0},\sigma_{p}^{2}\mathbf{I})\) with fixed standard deviation \(\sigma_{p}=1\). We train a two-layer CNN model defined in Section 2 with the ReLU activation function. The number of filters is set as \(m=10\). We use the default initialization method in PyTorch to initialize the CNN parameters and train the CNN with full-batch gradient descent with a learning rate of \(0.01\) for \(100\) iterations. We consider different dimensions \(d\) ranging from \(1000\) to \(20000\), and different signal strengths \(\|\mathbf{\mu}\|_{2}\) ranging from \(0\) to \(10\). Based on our results, for any dimension \(d\) and signal strength \(\mu\) setting we consider, our training setup can guarantee a training loss smaller than \(0.05\). After training, we estimate the test error for each case using \(1000\) test data points. We report the test error heat map with average results over \(10\) runs in Figure 2.

## 6 Related Work

**Sharpness Aware Minimization.**Foret et al. (2020), and Zheng et al. (2021) concurrently introduced methods to enhance generalization by minimizing the loss in the worst direction, perturbed from the current parameter. Kwon et al. (2021) introduced ASAM, a variant of SAM, designed to address parameter re-scaling. Subsequently, Liu et al. (2022b) presented LookSAM, a more computationally

Figure 2: (a) is a heatmap illustrating test error on synthetic data for various dimensions \(d\) and signal strengths \(\mathbf{\mu}\) when trained using Vanilla Gradient Descent. High test errors are represented in blue, while low test errors are shown in yellow. (b) displays a heatmap of test errors on the synthetic data under the same conditions as in (a), but trained using SAM instead with \(\tau=0.03\). The y-axis represents a normal scale with a range of \(1000\sim 21000\).

efficient alternative. Zhuang et al. (2022) highlighted that SAM did not consistently favor the flat minima and proposed GSAM to improve generalization by minimizing the surrogate. Recently, Zhao et al. (2022) showed that the SAM algorithm is related to the gradient regularization (GR) method when the loss is smooth and proposed an algorithm that can be viewed as a generalization of the SAM algorithm. Meng et al. (2023) further studied the mechanism of Per-Example Gradient Regularization (PEGR) on the CNN training and revealed that PEGR penalizes the variance of pattern learning.

**Benign Overfitting in Neural Networks.** Since the pioneering work by Bartlett et al. (2020) on benign overfitting in linear regression, there has been a surge of research studying benign overfitting in linear models, kernel methods, and neural networks. Li et al. (2021); Montanari and Zhong (2022) examined benign overfitting in random feature or neural tangent kernel models defined in two-layer neural networks. Chatterji and Long (2022) studied the excess risk of interpolating deep linear networks trained by gradient flow. Understanding benign overfitting in neural networks beyond the linear/kernel regime is much more challenging because of the non-convexity of the problem. Recently, Frei et al. (2022) studied benign overfitting in fully-connected two-layer neural networks with smoothed leaky ReLU activation. Cao et al. (2022) provided an analysis for learning two-layer convolutional neural networks (CNNs) with polynomial ReLU activation function (ReLU\({}^{q}\), \(q>2\)). Kou et al. (2023) further investigates the phenomenon of benign overfitting in learning two-layer ReLU CNNs. Kou et al. (2023) is most related to our paper. However, our work studied SGD rather than GD, which requires advanced techniques to control the update of coefficients at both batch-level and epoch-level. We also provide a novel analysis for SAM, which differs from the analysis of GD/SGD.

## 7 Conclusion

In this work, we rigorously analyze the training behavior of two-layer convolutional ReLU networks for both SGD and SAM. In particular, we precisely outlined the conditions under which benign overfitting can occur during SGD training, marking the first such finding for neural networks trained with mini-batch SGD. We also proved that SAM can lead to benign overfitting under circumstances that prompt harmful overfitting via SGD, which demonstrates the clear theoretical superiority of SAM over SGD. Our results provide a deeper comprehension of SAM, particularly when it comes to its utilization with non-smooth neural networks. An interesting future work is to consider other modern deep learning techniques, such as weight normalization, momentum, and weight decay, in our analysis.

## Acknowledgements

We thank the anonymous reviewers for their helpful comments. ZC, JZ, YK, and QG are supported in part by the National Science Foundation CAREER Award 1906169 and IIS-2008981, and the Sloan Research Fellowship. XC and CJH are supported in part by NSF under IIS-2008173, IIS-2048280, CISCO and Sony. The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies.

## References

* Ahn et al. (2022)Ahn, K., Bubeck, S., Chewi, S., Lee, Y. T., Suarez, F. and Zhang, Y. (2022). Learning threshold neurons via the" edge of stability". _arXiv preprint arXiv:2212.07469_.
* Allen-Zhu and Li (2020)Allen-Zhu, Z. and Li, Y. (2020). Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. _arXiv preprint arXiv:2012.09816_.
* Allen-Zhu and Li (2022)Allen-Zhu, Z. and Li, Y. (2022). Feature purification: How adversarial training performs robust deep learning. In _2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS)_. IEEE.
* Andriushchenko and Flammarion (2022)Andriushchenko, M. and Flammarion, N. (2022). Towards understanding sharpness-aware minimization. In _International Conference on Machine Learning_. PMLR.
* Andriushchenko et al. (2023)Andriushchenko, M., Varre, A. V., Pillaud-Vivien, L. and Flammarion, N. (2023). Sgd with large step sizes learns sparse features. In _International Conference on Machine Learning_. PMLR.
* Andriushchenko et al. (2021)Bahri, D., Mobahi, H. and Tay, Y. (2021). Sharpness-aware minimization improves language model generalization. _arXiv preprint arXiv:2110.08529_.
* Bartlett et al. (2022)Bartlett, P. L., Long, P. M. and Bousquet, O. (2022). The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. _arXiv preprint arXiv:2210.01513_.
* Bartlett et al. (2020)Bartlett, P. L., Long, P. M., Lugosi, G. and Tsigler, A. (2020). Benign overfitting in linear regression. _Proceedings of the National Academy of Sciences_.
* Behdin and Mazumder (2023)Behdin, K. and Mazumder, R. (2023). Sharpness-aware minimization: An implicit regularization perspective. _arXiv preprint arXiv:2302.11836_.
* Behdin et al. (2022)Behdin, K., Song, Q., Gupta, A., Durfee, D., Acharya, A., Keerthi, S. and Mazumder, R. (2022). Improved deep neural network generalization using m-sharpness-aware minimization. _arXiv preprint arXiv:2212.04343_.
* Cao et al. (2022)Cao, Y., Chen, Z., Belkin, M. and Gu, Q. (2022). Benign overfitting in two-layer convolutional neural networks. _Advances in neural information processing systems_**35** 25237-25250.
* Chatterji and Long (2021)Chatterji, N. S. and Long, P. M. (2021). Finite-sample analysis of interpolating linear classifiers in the overparameterized regime. _Journal of Machine Learning Research_**22** 129-1.
* Chatterji and Long (2022)Chatterji, N. S. and Long, P. M. (2022). Deep linear networks can benignly overfit when shallow ones do. _arXiv preprint arXiv:2209.09315_.
* Chen et al. (2021)Chen, X., Hsieh, C.-J. and Gong, B. (2021). When vision transformers outperform resnets without pre-training or strong data augmentations. _arXiv preprint arXiv:2106.01548_.
* Devroye et al. (2018)Devroye, L., Mehrabian, A. and Reddad, T. (2018). The total variation distance between high-dimensional gaussians. _arXiv preprint arXiv:1810.08693_.
* Foret et al. (2020)Foret, P., Kleiner, A., Mobahi, H. and Neyshabur, B. (2020). Sharpness-aware minimization for efficiently improving generalization. _arXiv preprint arXiv:2010.01412_.
* Frei et al. (2022)Frei, S., Chatterji, N. S. and Bartlett, P. (2022). Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In _Conference on Learning Theory_. PMLR.
* Jelassi and Li (2022)Jelassi, S. and Li, Y. (2022). Towards understanding how momentum improves generalization in deep learning. In _International Conference on Machine Learning_. PMLR.
* Kou et al. (2023)Kou, Y., Chen, Z., Chen, Y. and Gu, Q. (2023). Benign overfitting for two-layer relu networks. _arXiv preprint arXiv:2303.04145_.
* Kwon et al. (2021)Kwon, J., Kim, J., Park, H. and Choi, I. K. (2021). Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In _International Conference on Machine Learning_. PMLR.
* Li et al. (2019)Li, Y., Wei, C. and Ma, T. (2019). Towards explaining the regularization effect of initial large learning rate in training neural networks. In _Advances in Neural Information Processing Systems_.
* Li et al. (2021a)Li, Z., Wang, T. and Arora, S. (2021a). What happens after sgd reaches zero loss?-a mathematical framework. _arXiv preprint arXiv:2110.06914_.
* Li et al. (2021b)Li, Z., Zhou, Z.-H. and Gretton, A. (2021b). Towards an understanding of benign overfitting in neural networks. _arXiv preprint arXiv:2106.03212_.
* Liu et al. (2022a)Liu, Y., Mai, S., Chen, X., Hsieh, C.-J. and You, Y. (2022a). Towards efficient and scalable sharpness-aware minimization. In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_.
* Liu et al. (2022b)Liu, Y., Mai, S., Chen, X., Hsieh, C.-J. and You, Y. (2022b). Towards efficient and scalable sharpness-aware minimization. In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_.
* Liu et al. (2022)Meng, X., Cao, Y. and Zou, D. (2023). Per-example gradient regularization improves learning signals from noisy data. _arXiv preprint arXiv:2303.17940_.
* Montanari and Zhong (2022)Montanari, A. and Zhong, Y. (2022). The interpolation phase transition in neural networks: Memorization and generalization under lazy training. _The Annals of Statistics_**50** 2816-2847.
* Olshausen and Field (1997)Olshausen, B. A. and Field, D. J. (1997). Sparse coding with an overcomplete basis set: A strategy employed by v1? _Vision research_**37** 3311-3325.
* Shen et al. (2022)Shen, R., Bubeck, S. and Gunasekar, S. (2022). Data augmentation as feature manipulation: a story of desert cows and grass cows. _arXiv preprint arXiv:2203.01572_.
* Vershynin (2018)Vershynin, R. (2018). _High-Dimensional Probability: An Introduction with Applications in Data Science_. Cambridge Series in Statistical and Probabilistic Mathematics, Cambridge University Press.
* Wen et al. (2022)Wen, K., Ma, T. and Li, Z. (2022). How does sharpness-aware minimization minimize sharpness? _arXiv preprint arXiv:2211.05729_.
* Zhang et al. (2021)Zhang, C., Bengio, S., Hardt, M., Recht, B. and Vinyals, O. (2021). Understanding deep learning (still) requires rethinking generalization. _Communications of the ACM_**64** 107-115.
* Zhao et al. (2022)Zhao, Y., Zhang, H. and Hu, X. (2022). Penalizing gradient norm for efficiently improving generalization in deep learning. In _International Conference on Machine Learning_. PMLR.
* Zheng et al. (2021)Zheng, Y., Zhang, R. and Mao, Y. (2021). Regularizing neural networks via adversarial model perturbation. In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_.
* Zhuang et al. (2022)Zhuang, J., Gong, B., Yuan, L., Cui, Y., Adam, H., Dvornek, N., Tatikonda, S., Duncan, J. and Liu, T. (2022). Surrogate gap minimization improves sharpness-aware training. _arXiv preprint arXiv:2203.08065_.

## Appendix A Additional Experiments

### Real Experiments on CIFAR

In this section, we provide the experiments on real data sets.

**Varying different starting points for SAM.** In Section 4, we show that the SAM algorithm can effectively prevent noise memorization and thus improve feature learning in the early stage of training. Is SAM also effective if we add the algorithm in the middle of the training process? We conduct experiments on the ImageNet dataset with ResNet50. We choose the batch size as \(1024\), and the model is trained for \(90\) epochs with the best learning rate in grid search \(\{0.01,0.03,0.1,0.3\}\). The learning rate schedule is \(10\)k steps linear warmup, then cosine decay. As shown in Table 1, the earlier SAM is introduced, the more pronounced its effectiveness becomes.

**SAM with additive noises.** Here, we conduct experiments on the CIFAR dataset with WRN-16-8. We add Gaussian random noises to the image data with variance \(\{0.1,0.3,1\}\). We choose the batch size as \(128\) and train the model over \(200\) epochs using a learning rate of \(0.1\), a momentum of \(0.9\), and a weight decay of \(5e-4\). The SAM hyperparameter is chosen as \(\tau=2.0\). As we can see from Table 2, SAM can consistently prevent noise learning and get better performance, compared to the SGD, which varies from different additive noise levels.

### Discussion on the Stochastic Gradient Descent

Here, we include additional empirical results to show how the batch size and step size influence the transition phase of the benign/harmful regimes.

Our synthetic experiment is only performed with gradient descent instead of SGD in Figure 2. We add an experiment of SGD with a mini-batch size of 10 on the synthetic data. If we compare Figure 3 and Figure 2, the comparison results should remain the same, and both support our main Theorems 3.2 and 4.1.

In Figure 4, we also conducted an extended study on the learning rate to check whether tuning the learning rate can get better generalization performance and achieve SAM's performance. Specifically, we experimented with learning rates of \(0.001,0.01,0.1\), and \(1\) under the same conditions described

\begin{table}
\begin{tabular}{l c c c c} \hline \hline \(\tau\) & **10\%** & **30\%** & **50\%** & **70\%** & **90\%** \\ \hline
0.01 & 76.9 & 76.9 & 76.9 & 76.7 & 76.7 \\
0.02 & 77.1 & 77.0 & 76.9 & 76.8 & 76.6 \\
0.05 & 76.2 & 76.4 & 76.3 & 76.2 & \\ \hline \hline \end{tabular}
\end{table}
Table 1: Top-1 accuracy (%) of ResNet-50 on the ImageNet dataset when we vary the starting point of using the SAM update rule, baseline result is 76.4%.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline
**Model** & **Noise** & **Dataset** & **Optimizer** & **Accuracy** \\ \hline WRN-16-8 & - & CIFAR-10 & SGD & 96.69 \\ WRN-16-8 & - & CIFAR-10 & SAM & 97.19 \\ \hline WRN-16-8 & \(\mathcal{N}(0,0.1)\) & CIFAR-10 & SGD & 95.87 \\ WRN-16-8 & \(\mathcal{N}(0,0.1)\) & CIFAR-10 & SAM & 96.57 \\ \hline WRN-16-8 & \(\mathcal{N}(0,0.3)\) & CIFAR-10 & SGD & 92.40 \\ WRN-16-8 & \(\mathcal{N}(0,0.3)\) & CIFAR-10 & SAM & 93.37 \\ \hline WRN-16-8 & \(\mathcal{N}(0,1)\) & CIFAR-10 & SGD & 79.50 \\ WRN-16-8 & \(\mathcal{N}(0,1)\) & CIFAR-10 & SAM & 80.37 \\ \hline WRN-16-8 & - & CIFAR-100 & SGD & 81.93 \\ WRN-16-8 & - & CIFAR-100 & SAM & 83.68 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Top-1 accuracy (%) of wide ResNet on the CIFAR datasets when adding different levels of Gaussian noise.

in Section 5. The results show that for all learning rates, the harmful and benign overfitting patterns are quite similar and consistent with our Theorem 3.2.

We observed that larger learning rates, such as \(0.1\) and \(1\), can improve SGD's generalization performance. The benign overfitting region is enlarged for learning rates of \(0.1\) and \(1\) when contrasted with \(0.01\). This trend resonates with recent findings that large LR helps SGD find better minima (Li et al., 2019, 2021; Ahn et al., 2022; Andriushchenko et al., 2023). Importantly, even with this expansion, the benign overfitting region remains smaller than what is empirically observed with SAM. Our conclusion is that while SGD, with a more considerable learning rate, exhibits improved generalization, it still falls short of matching SAM's performance.

## Appendix B Preliminary Lemmas

**Lemma B.1** (Lemma B.4 in Kou et al. (2023)).: _Suppose that \(\delta>0\) and \(d=\Omega(\log(6n/\delta))\). Then with probability at least \(1-\delta\),_

\[\sigma_{p}^{2}d/2\leq\|\mathbf{\xi}_{i}\|_{2}^{2}\leq 3\sigma_{p}^{2}d /2,\] \[|\langle\mathbf{\xi}_{i},\mathbf{\xi}_{i^{\prime}}\rangle|\leq 2\sigma_{p} ^{2}\cdot\sqrt{d\log(6n^{2}/\delta)},\] \[|\langle\mathbf{\xi}_{i},\mathbf{\mu}\rangle|\leq\|\mathbf{\mu}\|_{2}\sigma_{ p}\cdot\sqrt{2\log(6n/\delta)}\]

_for all \(i,i^{\prime}\in[n]\)._

**Lemma B.2** (Lemma B.5 in Kou et al. (2023)).: _Suppose that \(d=\Omega(\log(mn/\delta))\), \(m=\Omega(\log(1/\delta))\). Then with probability at least \(1-\delta\),_

\[\sigma_{0}^{2}d/2\leq\|\mathbf{w}_{j,r}^{(0,0)}\|_{2}^{2}\leq 3 \sigma_{0}^{2}d/2,\] \[|\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\mu}\rangle|\leq\sqrt{2\log( 12m/\delta)}\cdot\sigma_{0}\|\mathbf{\mu}\|_{2},\] \[|\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle|\leq 2\sqrt{ \log(12mn/\delta)}\cdot\sigma_{0}\sigma_{p}\sqrt{d}\]

_for all \(r\in[m]\), \(j\in\{\pm 1\}\) and \(i\in[n]\). Moreover,_

\[\sigma_{0}\|\mathbf{\mu}\|_{2}/2\leq\max_{r\in[m]}j\cdot\langle\mathbf{ w}_{j,r}^{(0,0)},\mathbf{\mu}\rangle\leq\sqrt{2\log(12m/\delta)}\cdot\sigma_{0}\| \mathbf{\mu}\|_{2},\] \[\sigma_{0}\sigma_{p}\sqrt{d}/4\leq\max_{r\in[m]}j\cdot\langle \mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle\leq 2\sqrt{\log(12mn/\delta)}\cdot\sigma_{0} \sigma_{p}\sqrt{d}\]

_for all \(j\in\{\pm 1\}\) and \(i\in[n]\)._

**Lemma B.3**.: _Let \(S_{i}^{(t,b)}\) denote \(\{r:\langle\mathbf{w}_{y_{i},r}^{(t,b)},\mathbf{\xi}_{i}\rangle>\sigma_{0}\sigma_{ p}\sqrt{d}/\sqrt{2}\rangle\}\). Suppose that \(\delta>0\) and \(m\geq 50\log(2n/\delta)\). Then with probability at least \(1-\delta\),_

\[|S_{i}^{(0,0)}|\geq 0.8\Phi(-1)m,\,\forall i\in[n].\]

Figure 3: (a) is a heatmap illustrating test error on synthetic data for various dimensions \(d\) and signal strengths \(\mathbf{\mu}\) when trained using stochastic gradient descent. High test errors are represented in blue, while low test errors are shown in yellow. (b) displays a heatmap of test errors on the synthetic data under the same conditions as in (a), but trained using SAM instead. The y-axis represents a normal scale with a range of 1000-21000.

Proof of Lemma b.3.: Since \(\langle\mathbf{w}_{y_{i},r}^{(0,0)},\mathbf{\xi}_{i}\rangle\sim\mathcal{N}(0,\sigma_{0 }^{2}\|\mathbf{\xi}_{i}\|_{2}^{2})\), we have

\[P(\langle\mathbf{w}_{y_{i},r}^{(0,0)},\mathbf{\xi}_{i}\rangle>\sigma_{0}\sigma_{p} \sqrt{d}/\sqrt{2})\geq P(\langle\mathbf{w}_{y_{i},r}^{(0,0)},\mathbf{\xi}_{i} \rangle>\sigma_{0}\|\mathbf{\xi}_{i}\|_{2})=\Phi(-1),\]

where \(\Phi(\cdot)\) is CDF of the standard normal distribution. Note that \(|S_{i}^{(0,0)}|=\sum_{r=1}^{m}\mathds{1}[\langle\mathbf{w}_{y_{i},r}^{(0,0)}, \mathbf{\xi}_{i}\rangle>\sigma_{0}\sigma_{p}\sqrt{d}/\sqrt{2}]\) and \(P\big{(}\langle\mathbf{w}_{y_{i},r}^{(0,0)},\mathbf{\xi}_{i}\rangle>\sigma_{0} \sigma_{p}\sqrt{d}/\sqrt{2}\big{)}\geq\Phi(-1)\), then by Hoeffding's inequality, with probability at least \(1-\delta/n\), we have

\[\frac{|S_{i}^{(0,0)}|}{m}\geq\Phi(-1)-\sqrt{\frac{\log(2n/\delta)}{2m}}.\]

Therefore, as long as \(0.2\sqrt{m}\Phi(-1)\geq\sqrt{\frac{\log(2n/\delta)}{2}}\), by applying union bound, with probability at least \(1-\delta\), we have

\[|S_{i}^{(0)}|\geq 0.8\Phi(-1)m,\;\forall i\in[n].\]

**Lemma B.4**.: _Let \(S_{j,r}^{(t,b)}\) denote \(\{i\in[n]:y_{i}=j,\;\langle\mathbf{w}_{y_{i},r}^{(t,b)},\mathbf{\xi}_{i}\rangle> \sigma_{0}\sigma_{p}\sqrt{d}/\sqrt{2}\}\). Suppose that \(\delta>0\) and \(n\geq 32\log(4m/\delta)\). Then with probability at least \(1-\delta\),_

\[|S_{j,r}^{(0)}|\geq n\Phi(-1)/4,\;\forall j\in\{\pm 1\},r\in[m].\]

Figure 4: (a) is a heatmap illustrating test error on synthetic data for various dimensions \(d\) and signal strengths \(\mathbf{\mu}\) when trained using learning rate \(0.001\). High test errors are represented in blue, while low test errors are shown in yellow. (b)(c)(d) displays a heatmap of test errors on the synthetic data under the same conditions as in (a), but trained using GD with different learning rates \(0.01,0.1,1\). The y-axis represents a normal scale with a range of 1000-21000.

Proof of Lemma b.4.: Since \(\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle\sim\mathcal{N}(0,\sigma_{0}^{2 }\|\mathbf{\xi}_{i}\|_{2}^{2})\), we have

\[P(\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle>\sigma_{0}\sigma_{p} \sqrt{d}/\sqrt{2})\geq P(\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle> \sigma_{0}\|\mathbf{\xi}_{i}\|_{2})=\Phi(-1),\]

where \(\Phi(\cdot)\) is CDF of the standard normal distribution.

Note that \(|S_{j,r}^{(0,0)}|=\sum_{i=1}^{n}\mathds{1}[y_{i}=j]\mathds{1}[\langle\mathbf{w }_{j,r}^{(0)},\mathbf{\xi}_{i}\rangle>\sigma_{0}\sigma_{p}\sqrt{d}/\sqrt{2}]\) and \(\mathbb{P}(y_{i}=j,\langle\mathbf{w}_{j,r}^{(0)},\mathbf{\xi}_{i}\rangle>\sigma_{ 0}\sigma_{p}\sqrt{d}/\sqrt{2})\geq\Phi(-1)/2\), then by Hoeffding's inequality, with probability at least \(1-\delta/2m\), we have

\[\frac{|S_{j,r}^{(0)}|}{n}\geq\Phi(-1)/2+\sqrt{\frac{\log(4m/\delta)}{2n}}.\]

Therefore, as long as \(\Phi(-1)/4\geq\sqrt{\frac{\log(4m/\delta)}{2n}}\), by applying union bound, we have with probability at least \(1-\delta\),

\[|S_{j,r}^{(0)}|\geq n\Phi(-1)/4,\,\forall j\in\{\pm 1\},r\in[m].\]

**Lemma B.5** (Lemma B.3 in Kou et al. (2023)).: _For \(|S_{+}\cap S_{y}|\) and \(|S_{-}\cap S_{y}|\) where \(y\in\{\pm 1\}\), it holds with probability at least \(1-\delta(\delta>0)\) that_

\[\Big{|}|S_{+}\cap S_{y}|-\frac{(1-p)n}{2}\Big{|}\leq\sqrt{\frac{n}{2}\log\Big{(} \frac{8}{\delta}\Big{)}},\Big{|}|S_{-}\cap S_{y}|-\frac{pn}{2}\Big{|}\leq\sqrt {\frac{n}{2}\log\Big{(}\frac{8}{\delta}\Big{)}},\forall\,y\in\{\pm 1\}.\]

**Lemma B.6**.: _It holds with probability at least \(1-\delta\), for all \(T\in[\frac{\log(2T^{*}/\delta)}{c_{3}^{2}},T^{*}]\) and \(y\in\{\pm 1\}\), there exist at least \(c_{3}\cdot T\) epochs among \([0,T]\), such that at least \(c_{4}\cdot H\) batches in these epochs, satisfy_

\[|S_{+}\cap S_{y}\cap\mathcal{I}_{t,b}|\in\bigg{[}\frac{B}{4},\frac{3B}{4} \bigg{]}. \tag{16}\]

Proof.: Let

\[\mathcal{E}_{1,t} :=\{\text{In epoch $t$, there are at least $c_{2}\cdot\frac{n}{B}$ batches such that \ref{eq:t\[\mathbb{P}[\mathds{1}(\mathcal{E}_{1,t,c_{t}H-1})=l_{c_{1}H-1}|\, \mathds{1}(\mathcal{E}_{1,t,0})=l_{0},\cdots,\mathds{1}(\mathcal{E}_{1,t,c_{1}H-2 })=l_{c_{1}H-2}]\] \[\leq\sum_{i=0}^{c_{1}c_{2}H}C_{c_{1}H}^{i}(1-2c_{2})^{c_{1}H-i}\] \[\leq c_{1}c_{2}H\cdot(2c_{2})^{c_{1}c_{2}H}(1-2c_{2})^{c_{1}H-c_{1 }c_{2}H}.\]

Choose \(H_{0}\) such that \(c_{1}c_{2}H_{0}\cdot(2c_{2})^{c_{1}c_{2}H_{0}}(1-2c_{2})^{c_{1}H_{0}-c_{1}c_{2} H_{0}}=1-2c_{3}\), then as long as \(H\geq H_{0}\), with probability \(c_{3}\), there are at least \(c_{1}c_{2}H\) batches in first \(c_{1}H\) batches such that 16 holds. Then \(\mathbb{P}[\mathcal{E}_{1,t}]\geq 2c_{3}\). Therefore,

\[\mathbb{P}(\sum_{t}\mathds{1}(\mathcal{E}_{t})-2Tc_{3}\leq-t)\leq\exp(-\frac{2 t^{2}}{T}).\]

Let \(T\geq\frac{\log(2T^{*}/\delta)}{2c_{3}^{2}}\). Then, with probability at least \(1-\delta/(2T^{*})\),

\[\sum_{t}\mathds{1}(\mathcal{E}_{1,t})\geq c_{3}T.\]

Let \(c_{4}=c_{1}c_{2}\). Thus, there are at least \(c_{3}T^{*}\) epochs, such that they have at least \(c_{4}H\) batches satisfying EquationB.6. This also holds for \(y=-1\). Taking a union bound to get the result.

## Appendix C Proofs for SGD

In this section, we prove the results for SGD. We first define some notations. Define \(H=n/B\) as the number of batches within an epoch. For any \(t_{1},t_{2}\) and \(b_{1},b_{2}\in\overline{[H]}\), we write \((t_{1},b_{1})\leq(t,b)\leq(t_{2},b_{2})\) to denote all iterations from \(t_{1}\)-th epoch's \(b_{1}\)-th batch (included) to \(t_{2}\)-th epoch's \(b_{2}\)-th batch (included). And the meanings change accordingly if we replace \(\leq\) with \(<\).

### Signal-noise Decomposition Coefficient Analysis

This part is dedicated to analyzing the update rule of Signal-noise Decomposition Coefficients. It is worth noting that

\[F_{j}(\mathbf{W},\mathbf{X})=\frac{1}{m}\sum_{r=1}^{m}\sum_{p=1}^{P}\sigma( \langle\mathbf{w}_{j,r},\mathbf{x}_{p}\rangle)=\frac{1}{m}\sum_{r=1}^{m} \sigma(\langle\mathbf{w}_{j,r},\widehat{y}\mathbf{\mu}\rangle)+(P-1)\sigma( \langle\mathbf{w}_{j,r},\mathbf{\xi}\rangle).\]

Let \(\mathcal{I}_{t,b}\) denote the set of indices of randomly chosen samples at epoch \(t\) batch \(b\), and \(|\mathcal{I}_{t,b}|=B\), then the update rule is:

\[\text{for }b\in\overline{[H]}\qquad\mathbf{w}_{j,r}^{(t,b+1)} =\mathbf{w}_{j,r}^{(t,b)}-\eta\cdot\nabla_{\mathbf{w}_{j,r}}L_{ \mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})\] \[=\mathbf{w}_{j,r}^{(t,b)}-\frac{\eta(P-1)}{Bm}\sum_{i\in \mathcal{I}_{t,b}}\ell_{i}^{\prime(t,b)}\cdot\sigma^{\prime}(\langle\mathbf{w }_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle)\cdot jy_{i}\mathbf{\xi}_{i}\] \[\qquad-\frac{\eta}{Bm}\sum_{i\in\mathcal{I}_{t,b}}\ell_{i}^{ \prime(t,b)}\cdot\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\widehat{y}_ {i}\mathbf{\mu}\rangle)\cdot jy_{i}\widehat{y}_{i}\mathbf{\mu},\] \[\text{and}\qquad\mathbf{w}_{j,r}^{(t+1,0)} =\mathbf{w}_{j,r}^{(t,H)}. \tag{17}\]

#### c.1.1 Iterative Expression for Decomposition Coefficient Analysis

**Lemma C.1**.: _The coefficients \(\gamma_{j,r}^{(t,b)},\overline{\rho}_{j,r,i}^{(t,b)},\rho_{j,r,i}^{(t,b)}\) defined in Lemma 3.3 satisfy the following iterative equations:_

\[\gamma_{j,r}^{(0,0)},\overline{\rho}_{j,r,i}^{(0,0)},\rho_{j,r,i} ^{(0,0)}=0, \tag{18}\] \[\gamma_{j,r}^{(t,b+1)}=\gamma_{j,r}^{(t,b)}-\frac{\eta}{Bm}\cdot \left[\sum_{i\in\mathcal{I}_{t,b}\cap\mathcal{S}_{+}}\ell_{i}^{\prime(t,b)} \sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\widehat{y}_{i}\cdot\mathbf{\mu} \rangle)\right.\]\[\alpha :=4\log(T^{*}), \tag{22}\] \[\beta :=2\max_{i,j,r}\{|\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\mu}\rangle|,(P- 1)|\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle|\},\] (23) \[\mathrm{SNR} :=\frac{\|\mathbf{\mu}\|_{2}}{(P-1)\sigma_{p}\sqrt{d}}. \tag{24}\]

By Lemma B.2 and Condition 3.1, \(\beta\) can be bounded as

\[\beta=2\max_{i,j,r}\{|\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\mu}\rangle|,(P-1)| \langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle|\}\]\[\leq 2\max\{\sqrt{2\log(12m/\delta)}\cdot\sigma_{0}\|\mathbf{\mu}\|_{2},2 \sqrt{\log(12mn/\delta)}\cdot\sigma_{0}(P-1)\sigma_{p}\sqrt{d}\}\] \[=O\big{(}\sqrt{\log(mn/\delta)}\cdot\sigma_{0}(P-1)\sigma_{p}\sqrt {d}\big{)},\]

Then, by Condition 3.1, we have the following inequality:

\[\max\bigg{\{}\beta,\mathrm{SNR}\sqrt{\frac{32\log(6n/\delta)}{d}}n\alpha,5\sqrt {\frac{\log(6n^{2}/\delta)}{d}}n\alpha\bigg{\}}\leq\frac{1}{12}. \tag{25}\]

We first prove the following bounds for signal-noise decomposition coefficients.

**Proposition C.2**.: _Under Assumption 3.1, for \((0,0)\leq(t,b)\leq(T^{*},0)\), we have that_

\[\gamma_{j,r}^{(0,0)},\overline{\rho}_{j,r,i}^{(0,0)},\underline{ \rho}_{j,r,i}^{(0,0)}=0 \tag{26}\] \[0\leq\overline{\rho}_{j,r,i}^{(t,b)}\leq\alpha,\] (27) \[0\geq\underline{\rho}_{j,r,i}^{(t,b)}\geq-\beta-10\sqrt{\frac{ \log(6n^{2}/\delta)}{d}}n\alpha\geq-\alpha, \tag{28}\]

_and there exists a positive constant \(C^{\prime}\) such that_

\[-\frac{1}{12}\leq\gamma_{j,r}^{(t,b)}\leq C^{\prime}\widehat{\gamma}\alpha, \tag{29}\]

_for all \(r\in[m]\), \(j\in\{\pm 1\}\) and \(i\in[n]\), where \(\widehat{\gamma}:=n\cdot\mathrm{SNR}^{2}\)._

We will prove Proposition C.2 by induction. We first approximate the change of the inner product by corresponding decomposition coefficients when Proposition C.2 holds.

**Lemma C.3**.: _Under Assumption 3.1, suppose (27), (28) and (29) hold after \(b\)-th batch of \(t\)-th epoch. Then, for all \(r\in[m]\), \(j\in\{\pm 1\}\) and \(i\in[n]\),_

\[\big{|}\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(0,0)}, \mathbf{\mu}\rangle-j\cdot\gamma_{j,r}^{(t,b)}\big{|}\leq\mathrm{SNR}\sqrt{\frac{3 2\log(6n/\delta)}{d}}n\alpha, \tag{30}\] \[\big{|}\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(0,0)}, \mathbf{\xi}_{i}\rangle-\frac{1}{P-1}\underline{\rho}_{j,r,i}^{(t,b)}\big{|}\leq \frac{5}{P-1}\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha,\,j\neq y_{i},\] (31) \[\big{|}\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(0,0)}, \mathbf{\xi}_{i}\rangle-\frac{1}{P-1}\overline{\rho}_{j,r,i}^{(t,b)}\big{|}\leq \frac{5}{P-1}\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha,\,j=y_{i}. \tag{32}\]

Proof of Lemma c.3.: First, for any time \((t,b)\geq(0,0)\), we have from the following decomposition by definitions,

\[\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(0,0)},\mathbf{\mu }\rangle=j\cdot\gamma_{j,r}^{(t,b)}+\frac{1}{P-1}\sum_{i^{\prime}=1}^{n} \overline{\rho}_{j,r,i^{\prime}}^{(t,b)}\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2} \cdot\langle\mathbf{\xi}_{i^{\prime}},\mathbf{\mu}\rangle\] \[\qquad\qquad+\frac{1}{P-1}\sum_{i^{\prime}=1}^{n}\underline{\rho }_{j,r,i^{\prime}}^{(t,b)}\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot\langle\mathbf{ \xi}_{i^{\prime}},\mathbf{\mu}\rangle.\]

According to Lemma B.1, we have

\[\Bigg{|}\frac{1}{P-1}\sum_{i^{\prime}=1}^{n}\overline{\rho}_{j, r,i^{\prime}}^{(t,b)}\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot\langle\mathbf{\xi}_{i^{ \prime}},\mathbf{\mu}\rangle+\frac{1}{P-1}\sum_{i^{\prime}=1}^{n}\underline{\rho} _{j,r,i^{\prime}}^{(t,b)}\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot\langle\mathbf{ \xi}_{i^{\prime}},\mathbf{\mu}\rangle\Bigg{|}\] \[\leq\frac{1}{P-1}\sum_{i^{\prime}=1}^{n}|\overline{\rho}_{j,r,i^ {\prime}}^{(t,b)}\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot|\langle\mathbf{\xi}_{i^{ \prime}},\mathbf{\mu}\rangle|+\frac{1}{P-1}\sum_{i^{\prime}=1}^{n}|\underline{\rho }_{j,r,i^{\prime}}^{(t,b)}\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot|\langle \mathbf{\xi}_{i^{\prime}},\mathbf{\mu}\rangle|\] \[\leq\frac{2\|\mathbf{\mu}\|_{2}\sqrt{2\log(6n/\delta)}}{(P-1)\sigma_{ p}d}\bigg{(}\sum_{i^{\prime}=1}^{n}|\overline{\rho}_{j,r,i^{\prime}}^{(t,b)}|+ \sum_{i^{\prime}=1}^{n}|\underline{\rho}_{j,r,i^{\prime}}^{(t,b)}|\bigg{)}\]\[=\mathrm{SNR}\sqrt{\frac{8\log(6n/\delta)}{d}}\bigg{(}\sum_{i^{\prime}= 1}^{n}|\overline{\rho}_{j,r,i^{\prime}}^{(t,b)}|+\sum_{i^{\prime}=1}^{n}|\underline {\rho}_{j,r,i^{\prime}}^{(t,b)}|\bigg{)}\] \[\leq\mathrm{SNR}\sqrt{\frac{32\log(6n/\delta)}{d}}n\alpha,\]

where the first inequality is by triangle inequality, the second inequality is by Lemma B.1, the equality is by \(\mathrm{SNR}=\|\mathbf{\mu}\|_{2}/((P-1)\sigma_{p}\sqrt{d})\), and the last inequality is by (27), (28). It follows that

\[\big{|}\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(0,0)},\mathbf{\mu} \rangle-j\cdot\gamma_{j,r}^{(t,b)}\big{|}\leq\mathrm{SNR}\sqrt{\frac{32\log(6n /\delta)}{d}}n\alpha.\]

Then, for \(j\neq y_{i}\) and any \(t\geq 0\), we have

\[\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi} _{i}\rangle\] \[=j\cdot\gamma_{j,r}^{(t,b)}\|\mathbf{\mu}\|_{2}^{-2}\cdot\langle\mathbf{ \mu},\mathbf{\xi}_{i}\rangle+\frac{1}{P-1}\sum_{i^{\prime}=1}^{n}\overline{\rho}_{ j,r,i^{\prime}}^{(t,b)}\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot\langle\mathbf{\xi}_{i^{ \prime}},\mathbf{\xi}_{i}\rangle\] \[\qquad+\frac{1}{P-1}\sum_{i^{\prime}=1}^{n}\underline{\rho}_{j,r, i^{\prime}}^{(t,b)}\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot\langle\mathbf{\xi}_{i^{ \prime}},\mathbf{\xi}_{i}\rangle\] \[=j\cdot\gamma_{j,r}^{(t,b)}\|\mathbf{\mu}\|_{2}^{-2}\cdot\langle\mathbf{ \mu},\mathbf{\xi}_{i}\rangle+\frac{1}{P-1}\sum_{i^{\prime}=1}^{n}\underline{\rho}_ {j,r,i^{\prime}}^{(t,b)}\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot\langle\mathbf{\xi} _{i^{\prime}},\mathbf{\xi}_{i}\rangle\] \[=\frac{1}{P-1}\underline{\rho}_{j,r,i}^{(t,b)}+j\cdot\gamma_{j,r} ^{(t,b)}\|\mathbf{\mu}\|_{2}^{-2}\cdot\langle\mathbf{\mu},\mathbf{\xi}_{i}\rangle+\frac{1} {P-1}\sum_{i^{\prime}\neq i}\underline{\rho}_{j,r,i^{\prime}}^{(t,b)}\|\mathbf{\xi} _{i^{\prime}}\|_{2}^{-2}\cdot\langle\mathbf{\xi}_{i^{\prime}},\mathbf{\xi}_{i}\rangle,\]

where the second equality is due to \(\underline{\rho}_{j,r,i}^{(t,b)}=0\) for \(j\neq y_{i}\). Next, we have

\[\bigg{|}j\cdot\gamma_{j,r}^{(t,b)}\|\mathbf{\mu}\|_{2}^{-2}\cdot \langle\mathbf{\mu},\mathbf{\xi}_{i}\rangle+\frac{1}{P-1}\sum_{i^{\prime}\neq i}\underline {\rho}_{j,r,i^{\prime}}^{(t,b)}\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot\langle \mathbf{\xi}_{i^{\prime}},\mathbf{\xi}_{i}\rangle\bigg{|}\] \[\leq|\gamma_{j,r}^{(t,b)}\|\mathbf{\mu}\|_{2}^{-2}\cdot|\langle\mathbf{ \mu},\mathbf{\xi}_{i}\rangle|+\frac{1}{P-1}\sum_{i^{\prime}\neq i}|\underline{ \rho}_{j,r,i^{\prime}}^{(t,b)}\|\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot|\langle \mathbf{\xi}_{i^{\prime}},\mathbf{\xi}_{i}\rangle|\] \[\leq|\gamma_{j,r}^{(t,b)}\|\|\mathbf{\mu}\|_{2}^{-1}\sigma_{p}\sqrt{2 \log(6n/\delta)}+\frac{4}{P-1}\sqrt{\frac{\log(6n^{2}/\delta)}{d}}\sum_{i^{ \prime}\neq i}|\underline{\rho}_{j,r,i^{\prime}}^{(t,b)}|\] \[=\frac{\mathrm{SNR}^{-1}}{P-1}\sqrt{\frac{2\log(6n/\delta)}{d}}| \gamma_{j,r}^{(t,b)}|+\frac{4}{P-1}\sqrt{\frac{\log(6n^{2}/\delta)}{d}}\sum_{i^ {\prime}\neq i}|\underline{\rho}_{j,r,i^{\prime}}^{(t,b)}|\] \[\leq\frac{\mathrm{SNR}}{P-1}\sqrt{\frac{8C^{2}\log(6n/\delta)}{d }}n\alpha+\frac{4}{P-1}\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha\] \[\leq\frac{5}{P-1}\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha,\]

where the first inequality is by triangle inequality; the second inequality is by Lemma B.1; the equality is by \(\mathrm{SNR}=\|\mathbf{\mu}\|_{2}/\sigma_{p}\sqrt{d}\); the third inequality is by (28) and (29); the forth inequality is by \(\mathrm{SNR}\leq 1/\sqrt{8C^{\prime 2}}\). Therefore, for \(j\neq y_{i}\), we have

\[\big{|}\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i} \rangle-\frac{1}{P-1}\underline{\rho}_{j,r,i}^{(t,b)}\big{|}\leq\frac{5}{P-1} \sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha.\]

Similarly, we have for \(y_{i}=j\) that

\[\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle\] \[=j\cdot\gamma_{j,r}^{(t,b)}\|\mathbf{\mu}\|_{2}^{-2}\cdot\langle\mathbf{ \mu},\mathbf{\xi}_{i}\rangle+\frac{1}{P-1}\sum_{i^{\prime}=1}^{n}\overline{\rho}_{ j,r,i^{\prime}}^{(t,b)}\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot\langle\mathbf{\xi}_{i^{ \prime}},\mathbf{\xi}_{i}\rangle\]\[\leq 2\max\{\langle\mathbf{w}_{j,r}^{(t,b)},y_{i}\boldsymbol{\mu} \rangle,(P-1)\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\xi}_{i}\rangle,0\}\] \[\leq 6\max\left\{\langle\mathbf{w}_{j,r}^{(0)},y_{i}\boldsymbol{ \mu}\rangle,(P-1)\langle\mathbf{w}_{j,r}^{(0)},\boldsymbol{\xi}_{i}\rangle, \mathrm{SNR}\sqrt{\frac{32\log(6n/\delta)}{d}}n\alpha,y_{i}j\gamma_{j,r}^{(t, b)},\right.\] \[\left.\hskip 56.905512pt5\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n \alpha+\underline{\rho}_{j,r,i}^{(t,b)}\right\}\] \[\leq 6\max\left\{\beta/2,\mathrm{SNR}\sqrt{\frac{32\log(6n/\delta) }{d}}n\alpha,-\gamma_{j,r}^{(t,b)},5\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n \alpha\right\}\] \[\leq 0.5,\]

where the second inequality is by (30), (31) and (32); the third inequality is due to the definition of \(\beta\) and \(\underline{\rho}_{j,r,i}^{(t,b)}<0\); the third inequality is by (25) and \(-\gamma_{j,r}^{(t,b)}\leq\frac{1}{12}\).

**Lemma C.5**.: _Under Condition 3.1, suppose (27), (28) and (29) hold at \(b\)-th batch of \(t\)-th epoch. Then, it holds that_

\[(P-1)\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\xi}_{i}\rangle \geq-0.25,\] \[(P-1)\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\xi}_{i}\rangle \leq(P-1)\sigma(\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\xi}_{i} \rangle)\leq(P-1)\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\xi}_{i}\rangle+0.25,\]

_for any \(i\in[n]\)._Proof of Lemma c.5.: According to (32) in Lemma C.3, we have

\[(P-1)\langle\mathbf{w}^{(t,b)}_{y_{i},r},\mathbf{\xi}_{i}\rangle \geq(P-1)\langle\mathbf{w}^{(0,0)}_{y_{i},r},\mathbf{\xi}_{i}\rangle+ \overline{\rho}^{(t,b)}_{y_{i},r,i}-5n\sqrt{\frac{\log(6n^{2}/\delta)}{d}}\alpha\] \[\geq-\beta-5n\sqrt{\frac{\log(6n^{2}/\delta)}{d}}\alpha\] \[\geq-0.25,\]

where the second inequality is due to \(\overline{\rho}^{(t,b)}_{y_{i},r,i}\geq 0\), the third inequality is due to \(\beta<1/8\) and \(5n\sqrt{\log(6n^{2}/\delta)/d}\cdot\alpha<1/8\) by Condition 3.1.

For the second equation, the first inequality holds naturally since \(z\leq\sigma(z)\). For the inequality, if \(\langle\mathbf{w}^{(t)}_{y_{i},r},\mathbf{\xi}_{i}\rangle\leq 0\), we have

\[(P-1)\sigma(\langle\mathbf{w}^{(t,b)}_{y_{i},r},\mathbf{\xi}_{i}\rangle)=0\leq(P-1 )\langle\mathbf{w}^{(t,b)}_{y_{i},r},\mathbf{\xi}_{i}\rangle+0.25.\]

And if \(\langle\mathbf{w}^{(t,b)}_{y_{i},r},\mathbf{\xi}_{i}\rangle>0\), we have

\[(P-1)\sigma(\langle\mathbf{w}^{(t,b)}_{y_{i},r},\mathbf{\xi}_{i}\rangle)=(P-1) \langle\mathbf{w}^{(t,b)}_{y_{i},r},\mathbf{\xi}_{i}\rangle<(P-1)\langle\mathbf{w}^ {(t,b)}_{y_{i},r},\mathbf{\xi}_{i}\rangle+0.25.\]

**Lemma C.6** (Lemma C.6 in Kou et al. (2023)).: _Let \(g(z)=\ell^{\prime}(z)=-1/(1+\exp(z))\), then for all \(z_{2}-c\geq z_{1}\geq-1\) where \(c\geq 0\) we have that_

\[\frac{\exp(c)}{4}\leq\frac{g(z_{1})}{g(z_{2})}\leq\exp(c).\]

**Lemma C.7**.: _For any iteration \(t\in[0,T^{*})\) and \(b,b_{1},b_{2}\in\overline{[H]}\), we have the following statements hold:_

1. \(\Big{|}\sum_{r=1}^{m}\big{[}\overline{\rho}^{(t,0)}_{y_{i},r,i}-\overline{\rho }^{(t,0)}_{y_{k},r,k}\big{]}-\sum_{r=1}^{m}\big{[}\overline{\rho}^{(t,b_{1})} _{y_{i},r,i}-\overline{\rho}^{(t,b_{2})}_{y_{k},r,k}\big{]}\Big{|}\leq 0.1\kappa.\)__
2. \(\langle\mathbf{w}^{(t,b)}_{y_{i},r},\mathbf{\xi}_{i}\rangle\geq\langle\mathbf{w}^ {(t,0)}_{y_{i},r},\mathbf{\xi}_{i}\rangle-\sigma_{0}\sigma_{p}\sqrt{d}/\sqrt{2}\,.\)__
3. _Let_ \(\widetilde{S}^{(t,b)}_{i}=\{r\in[m]:\langle\mathbf{w}^{(t,b)}_{y_{i},r},\mathbf{ \xi}_{i}\rangle>0\}\)_, then we have_ \[S^{(t,0)}_{i}\subseteq\widetilde{S}^{(t,b)}_{i}.\]
4. _Let_ \(\widetilde{S}^{(t,b)}_{j,r}=\{i\in[n]:y_{i}=j,\langle\mathbf{w}^{(t,b)}_{j,r}, \mathbf{\xi}_{i}\rangle>0\}\)_, then we have_ \[S^{(t,0)}_{j,r}\subseteq\widetilde{S}^{(t,b)}_{j,r}.\]

Proof.: For the first statement,

\[\Big{|}\sum_{r=1}^{m}\big{[}\overline{\rho}^{(t,0)}_{y_{i},r,i}- \overline{\rho}^{(t,0)}_{y_{k},r,k}\big{]}-\sum_{r=1}^{m}\big{[}\overline{\rho }^{(t,b_{1})}_{y_{i},r,i}-\overline{\rho}^{(t,b_{2})}_{y_{k},r,k}\big{]}\Big{|}\] \[\qquad\leq\frac{\eta(P-1)^{2}}{Bm}\max\Big{\{}|S^{(\widetilde{t} -1,b_{1})}_{i}||\ell^{(\widetilde{t}-1,b_{1})}_{i}|\cdot\|\mathbf{\xi}_{i}\|_{2}^{ 2},|S^{(\widetilde{t}-1,b_{2})}_{k}||\ell^{(\widetilde{t}-1,b_{2})}_{k}| \cdot\|\mathbf{\xi}_{k}\|_{2}^{2}\Big{\}}\] \[\qquad\leq\frac{\eta(P-1)^{2}}{B}\frac{3\sigma_{p}^{2}d}{2}\] \[\qquad\leq 0.1\kappa,\]

where the first inequality follows from the iterative update rule of \(\overline{\rho}^{(t,b)}_{j,r,i}\), the second inequality is due to Lemma B.2, and the last inequality is due to Condition 3.1.

For the second statement, recall that the stochastic gradient update rule is

\[\langle\mathbf{w}^{(t,b)}_{y_{i},r},\mathbf{\xi}_{i}\rangle=\langle\mathbf{w}^{(t, b-1)}_{y_{i},r},\mathbf{\xi}_{i}\rangle-\frac{\eta}{Bm}\cdot\sum_{i^{\prime}\in \mathcal{I}_{t,b-1}}\ell^{(t,b-1)}_{i^{\prime}}\cdot\sigma^{\prime}(\langle \mathbf{w}^{(t,b-1)}_{y_{i},r},\mathbf{\mu}\rangle)\cdot\langle y_{i^{\prime}}\mathbf{ \mu},\mathbf{\xi}_{i}\rangle y_{i^{\prime}}\]\[-\frac{\eta(P-1)}{Bm}\cdot\sum_{i^{\prime}\in\mathcal{I}_{i,k-1}/i}\ell_{i^{\prime}} ^{(t,b-1)}\cdot\sigma^{\prime}(\langle\mathbf{w}_{y_{i},r}^{(t,b-1)},\mathbf{\xi}_ {i^{\prime}}\rangle)\cdot\langle\mathbf{\xi}_{i^{\prime}},\mathbf{\xi}_{i}\rangle.\]

Therefore,

\[\langle\mathbf{w}_{y_{i},r}^{(t,b)},\mathbf{\xi}_{i}\rangle \geq\langle\mathbf{w}_{y_{i},r}^{(t,0)},\mathbf{\xi}_{i}\rangle-\frac {\eta}{Bm}\cdot n\cdot\|\mu\|_{2}\sigma_{p}\sqrt{2\log(6n/\delta)}-\frac{\eta( P-1)}{Bm}\cdot n\cdot 2\sigma_{p}^{2}\sqrt{d\log(6n^{2}/\delta)}\] \[\geq\langle\mathbf{w}_{y_{i},r}^{(t,0)},\mathbf{\xi}_{i}\rangle-\sigma _{0}\sigma_{p}\sqrt{d}/\sqrt{2},\]

where the first inequality is due to Lemma B.1, and the second inequality is due to Condition 3.1.

For the third statement. Let \(r^{*}\in S_{i}^{(t,0)}\), then

\[\langle\mathbf{w}_{y_{i},r^{*}}^{(t,b)},\mathbf{\xi}_{i}\rangle\geq\langle\mathbf{ w}_{y_{i},r^{*}}^{(t,0)},\mathbf{\xi}_{i}\rangle-\sigma_{0}\sigma_{p}\sqrt{d}/ \sqrt{2}>0,\]

where the first inequality is due to the second statement, and the second inequality is due to the definition of \(\widetilde{S}_{i}^{t,0}\). Therefore, \(r^{*}\in\widetilde{S}_{i}^{(t,b)}\) and \(S_{i}^{(t,b)}\subseteq\widetilde{S}_{i}^{(t,b)}\). The fourth statement can be obtained similarly.

**Lemma C.8**.: _Under Assumption 3.1, suppose (27), (28) and (29) hold for any iteration \((t^{\prime},b^{\prime})\leq(t,0)\). Then, the following conditions also hold for \(\forall t^{\prime}\leq t\) and \(\forall b^{\prime},b^{\prime}_{1},b^{\prime}_{2}\in[H]\):_

1. \(\sum_{r=1}^{m}\left[\overline{\rho}_{y_{i},r_{i}}^{(t^{\prime},0)}-\overline{ \rho}_{y_{k},r_{k}}^{(t^{\prime},0)}\right]\leq\kappa\) _for all_ \(i,k\in[n]\)_._
2. \(y_{i}\cdot f(\mathbf{W}^{(t^{\prime},b^{\prime}_{1})},\mathbf{x}_{i})-y_{k} \cdot f(\mathbf{W}^{(t^{\prime},b^{\prime}_{2})},\mathbf{x}_{k})\leq C_{1}\) _for all_ \(i,k\in[n]\)_._
3. \(\ell_{i}^{(t^{\prime},b^{\prime}_{1})}/\ell_{k}^{(t^{\prime},b^{\prime}_{2})} \leq C_{2}=\exp(C_{1})\) _for all_ \(i,k\in[n]\)_._
4. \(S_{i}^{(0,0)}\subseteq S_{i}^{(t^{\prime},0)}\)_, where_ \(S_{i}^{(t^{\prime},0)}:=\{r\in[m]:\langle\mathbf{w}_{y_{i},r}^{(t^{\prime},0)},\mathbf{\xi}_{i}\rangle>\sigma_{0}\sigma_{p}\sqrt{d}/\sqrt{2}\}\)_, and hence_ \(|S_{i}^{(t^{\prime},0)}|\geq 0.8m\Phi(-1)\) _for all_ \(i\in[n]\)_._
5. \(S_{j,r}^{(0,0)}\subseteq S_{j,r}^{(t^{\prime},0)}\) _, where_ \(S_{j,r}^{(t^{\prime},0)}:=\{i\in[n]:y_{i}=j,\langle\mathbf{w}_{j,r}^{(t^{ \prime},0)},\mathbf{\xi}_{i}\rangle>\sigma_{0}\sigma_{p}\sqrt{d}/\sqrt{2}\}\)_, and hence_ \(|S_{j,r}^{(t^{\prime},0)}|\geq\Phi(-1)n/4\) _for all_ \(j\in\{\pm 1\},r\in[m]\)_._

_Here we take \(\kappa\) and \(C_{1}\) as \(10\) and \(5\) respectively._

Proof of Lemma c.8.: We prove Lemma C.8 by induction. When \(t^{\prime}=0\), the fourth and fifth conditions hold naturally by Lemma B.3 and B.4.

For the first condition, since we have \(\overline{\rho}_{j,r,i}^{(0,0)}=0\) for any \(j,r,i\) according to (26), it is straightforward that \(\sum_{r=1}^{m}\left[\overline{\rho}_{y_{i},r,i}^{(0,0)}-\overline{\rho}_{y_{k },r,k}^{(0,0)}\right]=0\) for all \(i,k\in[n]\). So the first condition holds for \(t^{\prime}=0\).

For the second condition, we have

\[y_{i}\cdot f(\mathbf{W}^{(0,0)},\mathbf{x}_{i})-y_{k}\cdot f( \mathbf{W}^{(0,0)},\mathbf{x}_{k})\] \[=F_{y_{i}}(\mathbf{W}_{y_{i}}^{(0,0)},\mathbf{x}_{i})-F_{-y_{i}} (\mathbf{W}_{-y_{i}}^{(0,0)},\mathbf{x}_{i})+F_{-y_{k}}(\mathbf{W}_{-y_{k}}^{(0,0)},\mathbf{x}_{i})-F_{y_{k}}(\mathbf{W}_{y_{k}}^{(0,0)},\mathbf{x}_{i})\] \[\leq F_{y_{i}}(\mathbf{W}_{y_{i}}^{(0,0)},\mathbf{x}_{i})+F_{-y_{k }}(\mathbf{W}_{-y_{k}}^{(0,0)},\mathbf{x}_{i})\] \[=\frac{1}{m}\sum_{r=1}^{m}[\sigma(\langle\mathbf{w}_{y_{i},r}^{(0,0)},y_{i}\mathbf{\mu}\rangle)+(P-1)\sigma(\langle\mathbf{w}_{y_{i},r}^{(0,0)}, \mathbf{\xi}_{i}\rangle)]\] \[\qquad+\frac{1}{m}\sum_{r=1}^{m}[\sigma(\langle\mathbf{w}_{-y_{k },r}^{(0,0)},y_{k}\mathbf{\mu}\rangle)+(P-1)\sigma(\langle\mathbf{w}_{-y_{k},r}^{(0,0)},\mathbf{\xi}_{i}\rangle)]\] \[\leq 4\beta\leq 1/3\leq C_{1},\]

where the first inequality is by \(F_{j}(\mathbf{W}_{j}^{(0,0)},\mathbf{x}_{i})>0\), the second inequality is due to (23), and the third inequality is due to (25).

By Lemma C.6 and the second condition, the third condition can be obtained directly as

\[\frac{\ell_{i}^{\prime(0,0)}}{\ell_{k}^{\prime(0,0)}}\leq\exp\big{(}y_{k}\cdot f( \mathbf{W}^{(0,0)},\mathbf{x}_{k})-y_{i}\cdot f(\mathbf{W}^{(0,0)},\mathbf{x}_{ i})\big{)}\leq\exp(C_{1}).\]

Now suppose there exists \((\widetilde{t},\widetilde{b})\leq(t,b)\) such that these five conditions hold for any \((0,0)\leq(t^{\prime},b^{\prime})<(\widetilde{t},\widetilde{b})\). We aim to prove that these conditions also hold for \((t^{\prime},b^{\prime})=(\widetilde{t},\widetilde{b})\).

We first show that, for any \(0\leq t^{\prime}\leq t\) and \(0\leq b^{\prime}_{1},b^{\prime}_{2}\leq b\), \(y_{i}\cdot f(\mathbf{W}^{(t^{\prime},b^{\prime}_{1}},\mathbf{x}_{i})-y_{k} \cdot f(\mathbf{W}^{(t^{\prime},b^{\prime}_{2}}),\mathbf{x}_{k})\) can be approximated by \(\frac{1}{m}\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r_{i}}^{(t^{\prime},b^ {\prime}_{1})}-\overline{\rho}_{y_{k},r,k}^{(t^{\prime},b^{\prime}_{2})}\big{]}\) with a small constant approximation error. We begin by writing out

\[y_{i}\cdot f(\mathbf{W}^{(t^{\prime},b^{\prime}_{1}},\mathbf{x}_ {i})-y_{k}\cdot f(\mathbf{W}^{(t^{\prime},b^{\prime}_{2}},\mathbf{x}_{k})\] \[=y_{i}\sum_{j\in\{\pm 1\}}j\cdot F_{j}(\mathbf{W}^{(t^{\prime},b^{ \prime}_{1}},\mathbf{x}_{i})-y_{k}\sum_{j\in\{\pm 1\}}j\cdot F_{j}(\mathbf{W}^{(t^{ \prime},b^{\prime}_{2}}_{j},\mathbf{x}_{k})\] \[=F_{-y_{k}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{2}}_{-y_{k}}, \mathbf{x}_{k})-F_{-y_{i}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{1}}_{-y_{i}}, \mathbf{x}_{i})+F_{y_{i}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{1}}_{y_{i}}, \mathbf{x}_{i})-F_{y_{k}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{2}}_{y_{k}}, \mathbf{x}_{k})\] \[=F_{-y_{k}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{2}}_{-y_{k}}, \mathbf{x}_{k})-F_{-y_{i}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{1}}_{-y_{i}}, \mathbf{x}_{i})\] \[\qquad+\frac{1}{m}\sum_{r=1}^{m}[\sigma(\langle\mathbf{w}^{(t^{ \prime},b^{\prime}_{1}}_{y_{i},r},y_{i}\cdot\boldsymbol{\mu}\rangle)+(P-1) \sigma(\langle\mathbf{w}^{(t^{\prime},b^{\prime}_{1}}_{y_{i},r},\boldsymbol{ \xi}_{i}\rangle)]\] \[\qquad-\frac{1}{m}\sum_{r=1}^{m}[\sigma(\langle\mathbf{w}^{(t^{ \prime},b^{\prime}_{2}}_{y_{k},r},y_{k}\cdot\boldsymbol{\mu}\rangle)+(P-1) \sigma(\langle\mathbf{w}^{(t^{\prime},b^{\prime}_{2}}_{y_{k},r},\boldsymbol{ \xi}_{k}\rangle)]\] \[=\underbrace{F_{-y_{k}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{2}}_{ -y_{k}},\mathbf{x}_{k})-F_{-y_{i}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{1}}_{-y _{i}},\mathbf{x}_{i})}_{\text{I}_{1}}\] \[\qquad+\underbrace{\frac{1}{m}\sum_{r=1}^{m}[\sigma(\langle \mathbf{w}^{(t^{\prime},b^{\prime}_{1}}_{y_{i},r},y_{i}\cdot\boldsymbol{\mu} \rangle)-\sigma(\langle\mathbf{w}^{(t^{\prime},b^{\prime}_{2}}_{y_{k},r},y_{k }\cdot\boldsymbol{\mu}\rangle)]}_{\text{I}_{2}}\] \[\qquad+\underbrace{\frac{1}{m}\sum_{r=1}^{m}[(P-1)\sigma( \langle\mathbf{w}^{(t^{\prime},b^{\prime}_{1}}_{y_{i},r},\boldsymbol{\xi}_{i} \rangle)-(P-1)\sigma(\langle\mathbf{w}^{(t^{\prime},b^{\prime}_{2}}_{y_{k},r}, \boldsymbol{\xi}_{k}\rangle)]}_{\text{I}_{3}}, \tag{33}\]

where all the equalities are due to the network definition. Then we bound \(\text{I}_{1}\), \(\text{I}_{2}\) and \(\text{I}_{3}\).

For \(|\text{I}_{1}|\), we have the following upper bound by Lemma C.4:

\[|\text{I}_{1}| \leq|F_{-y_{k}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{2}}_{-y_{k}},\mathbf{x}_{k})|+|F_{-y_{i}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{1}}_{-y_{i}}, \mathbf{x}_{i})|\] \[=F_{-y_{k}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{2}}_{-y_{k}}, \mathbf{x}_{k})+F_{-y_{i}}(\mathbf{W}^{(t^{\prime},b^{\prime}_{1}}_{-y_{i}}, \mathbf{x}_{i})\] \[\leq 1. \tag{34}\]

For \(|\text{I}_{2}|\), we have the following upper bound:

\[|\text{I}_{2}| \leq\max\left\{\frac{1}{m}\sum_{r=1}^{m}\sigma(\langle\mathbf{w }^{(t^{\prime},b^{\prime}_{1}}_{y_{i},r},y_{i}\cdot\boldsymbol{\mu}\rangle), \frac{1}{m}\sum_{r=1}^{m}\sigma(\langle\mathbf{w}^{(t^{\prime},b^{\prime}_{2}}_{ y_{k},r},y_{k}\cdot\boldsymbol{\mu}\rangle)\right\}\] \[\leq 3\max\left\{|\langle\mathbf{w}^{(0,0)}_{y_{i},r},y_{i}\cdot \boldsymbol{\mu}\rangle|,|\langle\mathbf{w}^{(0,0)}_{y_{k},r},y_{k}\cdot \boldsymbol{\mu}\rangle|,\gamma^{(t^{\prime},b^{\prime}_{1}}_{j,r},\gamma^{(t^{ \prime},b^{\prime}_{2}}_{j,r},\text{SNR}\sqrt{\frac{32\log(6n/\delta)}{d}}n \alpha\right\}\] \[\leq 3\max\left\{\beta,C^{\prime}\widehat{\gamma}\alpha,\text{SNR} \sqrt{\frac{32\log(6n/\delta)}{d}}n\alpha\right\}\] \[\leq 0.25, \tag{35}\]where the second inequality is due to (30), the second inequality is due to the definition of \(\beta\) and (29), the third inequality is due to Condition 3.1 and (25).

For \(\mathrm{I}_{3}\), we have the following upper bound

\[\mathrm{I}_{3} =\frac{1}{m}\sum_{r=1}^{m}\big{[}(P-1)\sigma(\langle\mathbf{w}_{y_ {i},r}^{(t^{\prime},b_{1}^{\prime})},\boldsymbol{\xi}_{i}\rangle)-(P-1)\sigma( \langle\mathbf{w}_{y_{k},r}^{(t^{\prime},b_{2}^{\prime})},\boldsymbol{\xi}_{k }\rangle)\big{]}\] \[\leq\frac{1}{m}\sum_{r=1}^{m}\big{[}(P-1)\langle\mathbf{w}_{y_{i},r}^{(t^{\prime},b_{1}^{\prime})},\boldsymbol{\xi}_{i}\rangle-(P-1)\langle \mathbf{w}_{y_{k},r}^{(t^{\prime},b_{2}^{\prime})},\boldsymbol{\xi}_{k}\rangle \big{]}+0.25\] \[\leq\frac{1}{m}\sum_{r=1}^{m}\Big{[}\overline{\rho}_{y_{i},r,i}^ {(t^{\prime},b_{1}^{\prime})}-\overline{\rho}_{y_{k},r,k}^{(t^{\prime},b_{2}^ {\prime})}+10\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha\Big{]}+0.25\] \[\leq\frac{1}{m}\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^ {(t^{\prime},b_{1}^{\prime})}-\overline{\rho}_{y_{k},r,k}^{(t^{\prime},b_{2}^ {\prime})}\big{]}+0.5, \tag{36}\]

where the first inequality is due to Lemma C.5, the second inequality is due to Lemma C.3, the third inequality is due to \(5\sqrt{\log(6n^{2}/\delta)/d}n\alpha\leq 1/8\) according to Condition 3.1.

Similarly, we have the following lower bound

\[\mathrm{I}_{3} =\frac{1}{m}\sum_{r=1}^{m}\big{[}(P-1)\sigma(\langle\mathbf{w}_{y _{i},r}^{(t^{\prime},b_{1}^{\prime})},\boldsymbol{\xi}_{i}\rangle)-(P-1)\sigma (\langle\mathbf{w}_{y_{k},r}^{(t^{\prime},b_{2}^{\prime})},\boldsymbol{\xi}_{k }\rangle)\big{]}\] \[\geq\frac{1}{m}\sum_{r=1}^{m}\big{[}(P-1)\langle\mathbf{w}_{y_{i},r}^{(t^{\prime},b_{1}^{\prime})},\boldsymbol{\xi}_{i}\rangle-(P-1)\langle \mathbf{w}_{y_{k},r}^{(t^{\prime},b_{2}^{\prime})},\boldsymbol{\xi}_{k}\rangle \big{]}-0.25\] \[\geq\frac{1}{m}\sum_{r=1}^{m}\Big{[}\overline{\rho}_{y_{i},r,i}^ {(t^{\prime},b_{1}^{\prime})}-\overline{\rho}_{y_{k},r,k}^{(t^{\prime},b_{2}^ {\prime})}-10\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha\Big{]}-0.25\] \[\geq\frac{1}{m}\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^ {(t^{\prime},b_{1}^{\prime})}-\overline{\rho}_{y_{k},r,k}^{(t^{\prime},b_{2}^ {\prime})}\big{]}-0.5, \tag{37}\]

where the first inequality is due to Lemma C.5, the second inequality is due to Lemma C.3, the third inequality is due to \(5\sqrt{\log(6n^{2}/\delta)/d}n\alpha\leq 1/8\) according to Condition 3.1.

By plugging (34)-(36) into (33), we have

\[y_{i}\cdot f(\mathbf{W}^{(t^{\prime},b_{1}^{\prime})},\mathbf{x }_{i})-y_{k}\cdot f(\mathbf{W}^{(t^{\prime},b_{2}^{\prime})},\mathbf{x}_{k}) \leq|\mathrm{I}_{1}|+|\mathrm{I}_{2}|+\mathrm{I}_{3}\] \[\leq\frac{1}{m}\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^ {(t^{\prime},b_{1}^{\prime})}-\overline{\rho}_{y_{k},r,k}^{(t^{\prime},b_{2}^ {\prime})}\big{]}+1.75,\] \[y_{i}\cdot f(\mathbf{W}^{(t^{\prime},b_{1}^{\prime})},\mathbf{x }_{i})-y_{k}\cdot f(\mathbf{W}^{(t^{\prime},b_{2}^{\prime})},\mathbf{x}_{k}) \geq-|\mathrm{I}_{1}|-|\mathrm{I}_{2}|+\mathrm{I}_{3}\] \[\geq\frac{1}{m}\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^ {(t^{\prime},b_{1}^{\prime})}-\overline{\rho}_{y_{k},r,k}^{(t^{\prime},b_{2}^ {\prime})}\big{]}-1.75,\]

which is equivalent to

\[\left|y_{i}\cdot f(\mathbf{W}^{(t^{\prime},b_{1}^{\prime})},\mathbf{x}_{i})-y_ {k}\cdot f(\mathbf{W}^{(t^{\prime},b_{2}^{\prime})},\mathbf{x}_{k})-\frac{1}{m }\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{(t^{\prime},b_{1}^{\prime})} -\overline{\rho}_{y_{k},r,k}^{(t^{\prime},b_{2}^{\prime})}\big{]}\right|\leq 1.75. \tag{38}\]

Therefore, the second condition immediately follows from the first condition.

Then, we prove the first condition holds for \((\widetilde{t},\widetilde{b})\). Recall from Lemma C.1 that

\[\overline{\rho}_{j,r,i}^{(t,b+1)}=\overline{\rho}_{j,r,i}^{(t,b)}-\frac{\eta(P- 1)^{2}}{Bm}\cdot\ell_{i}^{(t,b)}\cdot\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{ (t,b)},\boldsymbol{\xi}_{i}\rangle)\cdot\|\boldsymbol{\xi}_{i}\|_{2}^{2}\cdot \mathds{1}(y_{i}=j)\,\mathds{1}(i\in\mathcal{I}_{t,b})\]

for all \(j\in\{\pm 1\},r\in[m],i\in[n],(0,0)\leq(t,b)<[T^{*},0]\). It follows that

\[\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{(t,b+1)}-\overline{\rho}_{y_{k },r,k}^{(t,b+1)}\big{]}\]\[=\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{(t,b)}-\overline{\rho}_{y_{k},r,k }^{(t,b)}\big{]}-\frac{\eta(P-1)^{2}}{Bm}\cdot\big{(}|\widetilde{S}_{i}^{(t,b)} |\ell_{i}^{\prime(t,b)}\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}\,\mathds{1}(i\in\mathcal{I }_{t,b})\]

for all \(i,k\in[n]\) and \(0\leq t\leq T^{*}\), \(b<H\).

If \(\widetilde{b}\in\{1,2,\cdots,H-1\}\), then the first statement for \((t^{\prime},b^{\prime})=(\widetilde{t},\widetilde{b})\) and for the last \((t^{\prime},b^{\prime})<(\widetilde{t},\widetilde{b})\) are the same. Otherwise, if \(\widetilde{b}=0\), we consider two separate cases: \(\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{(\widetilde{t}-1,0)}- \overline{\rho}_{y_{k},r,k}^{(\widetilde{t}-1,0)}\big{]}\leq 0.9\kappa\) and \(\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{(\widetilde{t}-1,0)}- \overline{\rho}_{y_{k},r,k}^{(\widetilde{t}-1,0)}\big{]}>0.9\kappa\).

When \(\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{(\widetilde{t}-1,0)}- \overline{\rho}_{y_{k},r,k}^{(\widetilde{t}-1,0)}\big{]}\leq 0.9\kappa\), we have

\[\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{(\widetilde{t} -1,0)}-\overline{\rho}_{y_{k},r,k}^{(\widetilde{t}-1,0)}\big{]}\] \[=\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{(\widetilde{t} -1,0)}-\overline{\rho}_{y_{k},r,k}^{(\widetilde{t}-1,0)}\big{]}-\frac{\eta(P-1) ^{2}}{Bm}\cdot\Big{(}|\widetilde{S}_{i}^{(\widetilde{t}-1,b_{i}^{(\widetilde{t }-1)})}|\ell_{i}^{(\widetilde{t}-1,b_{i}^{(\widetilde{t}-1)})}\cdot\|\mathbf{\xi}_ {i}\|_{2}^{2}\] \[-\big{|}\widetilde{S}_{k}^{(\widetilde{t}-1,b_{k}^{(\widetilde{t }-1)})}|\ell_{k}^{(\widetilde{t}-1,b_{k}^{(\widetilde{t}-1)})}\cdot\|\mathbf{\xi}_ {k}\|_{2}^{2}\Big{)}\] \[\leq\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{(\widetilde {t}-1,0)}-\overline{\rho}_{y_{k},r,k}^{(\widetilde{t}-1,0)}\big{]}-\frac{\eta(P -1)^{2}}{Bm}\cdot|\mathbf{\tilde{S}}_{i}^{(\widetilde{t}-1,b_{i}^{(\widetilde{t}-1 )})}|\ell_{i}^{(\widetilde{t}-1,b_{i}^{(\widetilde{t}-1)})}\cdot\|\mathbf{\xi}_{i} \|_{2}^{2}\] \[\leq\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{(\widetilde {t}-1,0)}-\overline{\rho}_{y_{k},r,k}^{(\widetilde{t}-1,0)}\big{]}+\frac{\eta(P -1)^{2}}{B}\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}\] \[\leq 0.9\kappa+0.1\kappa\] \[=\kappa,\]

where the first inequality is due to \(\ell_{i}^{(\widetilde{t}-1,b_{i}^{(\widetilde{t}-1)})}<0\); the second inequality is due to \(\big{|}S_{i}^{(\widetilde{t}-1,b_{i}^{(\widetilde{t}-1)})}\big{|}\leq m\) and \(-\ell_{i}^{\prime(\widetilde{t}-1,b_{i}^{(\widetilde{t}-1)})}<1\); the third inequality is due to Condition 3.1.

On the other hand, for when \(\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{(\widetilde{t}-1,0)}- \overline{\rho}_{y_{k},r,k}^{(\widetilde{t}-1,0)}\big{]}>0.9\kappa\), we have from the (38) that

\[y_{i}\cdot f(\mathbf{W}^{(\widetilde{t}-1,b_{i}^{(\widetilde{t}- 1)})},\mathbf{x}_{i})-y_{k}\cdot f(\mathbf{W}^{(\widetilde{t}-1,b_{k}^{( \widetilde{t}-1)})},\mathbf{x}_{k})\] \[\geq\frac{1}{m}\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^{( \widetilde{t}-1,b_{i}^{(\widetilde{t}-1)})}-\overline{\rho}_{y_{k},r,k}^{( \widetilde{t}-1,b_{k}^{(\widetilde{t}-1)})}\big{]}-1.75\] \[\geq\frac{1}{m}\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^ {(\widetilde{t}-1,0)}-\overline{\rho}_{y_{k},r,k}^{(\widetilde{t}-1,0)}\big{]}- 0.1\kappa-1.75\] \[\geq 0.9\kappa-0.1\kappa-0.54\kappa\] \[=0.26\kappa, \tag{39}\]

where the second inequality is due to \(\kappa=10\). Thus, according to Lemma C.6, we have

\[\frac{\ell_{i}^{\prime(\widetilde{t}-1,b_{i}^{(\widetilde{t}-1)})}}{\ell_{k}^{ \prime(\widetilde{t}-1,b_{k}^{(\widetilde{t}-1)})}}\leq\exp\big{(}y_{k}\cdot f (\mathbf{W}^{(\widetilde{t}-1,b_{k}^{(\widetilde{t}-1)})},\mathbf{x}_{k})-y_{i} \cdot f\big{(}\mathbf{W}^{(\widetilde{t}-1,b_{i}^{(\widetilde{t}-1)})},\mathbf{x} _{i})\big{)}\leq\exp(-0.26\kappa).\]

Since \(S_{i}^{(\widetilde{t}-1,0)}\subseteq\widetilde{S}_{i}^{(\widetilde{t}-1,b_{i}^{( \widetilde{t}-1)})}\), we have \(\left|\widetilde{S}_{k}^{(\widetilde{t}-1,b_{k}^{(\widetilde{t}-1)})}\right| \geq 0.8\Phi(-1)m\) according to the fourth condition. Also we have that \(|S_{i}^{(\widetilde{t}-1,b_{i}^{(\widetilde{t}-1)})}|\leq m\). It follows that

\[\frac{\big{|}S_{i}^{(\widetilde{t}-1,b_{i}^{(\widetilde{t}-1)})}\big{|}\ell_{i}^ {(\widetilde{t}-1,b_{i}^{(\widetilde{t}-1)})}\big{|}\ell_{k}^{(\widetilde{t}-1,b_{k }^{(\widetilde{t}-1)})}\big{|}\ell_{k}^{(\widetilde{t}-1,b_{k}^{(\widetilde{t}-1)}) }\leq\frac{\exp(-0.26\kappa)}{0.8\Phi(-1)}<0.8.\]According to Lemma B.1, under event \(\mathcal{E}_{\mathrm{prelim}}\), we have

\[\big{|}\|\mathbf{\xi}_{i}\|_{2}^{2}-d\cdot\sigma_{p}^{2}|=O\big{(}\sigma_{p}^{2}\cdot \sqrt{d\log(6n/\delta)}\big{)},\,\forall i\in[n].\]

Note that \(d=\Omega(\log(6n/\delta))\) from Condition 3.1, it follows that

\[|S_{i}^{(\widetilde{t},b_{i}^{(\widetilde{t}-1)})}|(-\ell_{i}^{(\widetilde{t},b _{i}^{(\widetilde{t}-1)})})\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}<|S_{k}^{(\widetilde{t },b_{k}^{(\widetilde{t}-1)})}|(-\ell_{k}^{(\widetilde{t},b_{k}^{(\widetilde{t} -1)})})\cdot\|\mathbf{\xi}_{k}\|_{2}^{2}.\]

Then we have

\[\sum_{r=1}^{m}\big{[}\overline{p}_{y_{i},r,i}^{(\widetilde{t},0)}-\overline{p} _{y_{k},r,k}^{(\widetilde{t},0)}\big{]}\leq\sum_{r=1}^{m}\big{[}\overline{p}_ {y_{i},r,i}^{(\widetilde{t}-1,0)}-\overline{p}_{y_{k},r,k}^{(\widetilde{t}-1, 0)}\big{]}\leq\kappa,\]

which completes the proof of the first hypothesis at iteration \((t^{\prime},b^{\prime})=(\widetilde{t},\widetilde{b})\). Next, by applying the approximation in (38), we are ready to verify the second hypothesis at iteration \((\widetilde{t},\widetilde{b})\). In fact, for any \((t^{\prime},b_{1}^{\prime}),(t^{\prime},b_{2}^{\prime})\leq(\widetilde{t}, \widetilde{b})\), we have

\[y_{i}\cdot f(\mathbf{W}^{(t^{\prime},b_{1}^{\prime})},\mathbf{x} _{i})-y_{k}\cdot f(\mathbf{W}^{(t^{\prime},b_{2}^{\prime})},\mathbf{x}_{k}) \leq\frac{1}{m}\sum_{r=1}^{m}\big{[}\overline{p}_{y_{i},r,i}^{(t^ {\prime},b_{1}^{\prime})}-\overline{p}_{y_{k},r,k}^{(t^{\prime},b_{2}^{ \prime})}\big{]}+1.75\] \[\leq\frac{1}{m}\sum_{r=1}^{m}\big{[}\overline{\rho}_{y_{i},r,i}^ {(t^{\prime},0)}-\overline{\rho}_{y_{k},r,k}^{(t^{\prime},0)}\big{]}+0.1\kappa +1.75\] \[\leq C_{1},\]

where the first inequality is by (38); the last inequality is by induction hypothesis and taking \(\kappa\) as 10 and \(C_{1}\) as 5.

And the third hypothesis directly follows by noting that, for any \((t^{\prime},b_{1}^{\prime}),(t^{\prime},b_{2}^{\prime})\leq(\widetilde{t}, \widetilde{b})\),

\[\frac{\ell_{i}^{(t^{\prime},b_{1}^{\prime})}}{\ell_{k}^{(t^{\prime},b_{2}^{ \prime})}}\leq\exp\big{(}y_{k}\cdot f(\mathbf{W}^{(t^{\prime},b_{1}^{\prime})},\mathbf{x}_{k})-y_{i}\cdot f(\mathbf{W}^{(t^{\prime},b_{2}^{\prime})}, \mathbf{x}_{i})\big{)}\leq\exp(C_{1})=C_{2}.\]

For the fourth hypothesis, If \(\widetilde{b}\in\{1,2,\cdots,H-1\}\), then the first statement for \((t^{\prime},b^{\prime})=(\widetilde{t},\widetilde{b})\) and for the last \((t^{\prime},b^{\prime})<(\widetilde{t},\widetilde{b})\) are the same. Otherwise, if \(\widetilde{b}=0\), according to the gradient descent rule, we have

\[\langle\mathbf{w}_{y_{i},r}^{(\widetilde{t},0)},\mathbf{\xi}_{i}\rangle =\langle\mathbf{w}_{y_{i},r}^{(\widetilde{t}-1,0)},\mathbf{\xi}_{i} \rangle-\frac{\eta}{Bm}\cdot\sum_{b^{\prime}=0}^{H-1}\sum_{i^{\prime}\in \mathcal{I}_{\widetilde{t}-1,b^{\prime}}}\ell_{i^{\prime}}^{(\widetilde{t}-1, \widetilde{b})}\cdot\sigma^{\prime}(\langle\mathbf{w}_{y_{i},r}^{(\widetilde{t }-1,b^{\prime})},y_{i^{\prime}}\mathbf{\mu}\rangle)\cdot\langle y_{i^{\prime}}\mathbf{ \mu},\mathbf{\xi}_{i}\rangle y_{i^{\prime}}\] \[\qquad-\frac{\eta(P-1)}{Bm}\cdot\sum_{b^{\prime}=0}^{H-1}\sum_{i^{ \prime}\in\mathcal{I}_{\widetilde{t}-1,b^{\prime}}}\ell_{i^{\prime}}^{( \widetilde{t}-1,b^{\prime})}\cdot\sigma^{\prime}(\langle\mathbf{w}_{y_{i},r}^{ (\widetilde{t}-1,\widetilde{b})},\mathbf{\xi}_{i^{\prime}}\rangle)\cdot\langle\bm {\xi}_{i^{\prime}},\mathbf{\xi}_{i}\rangle\] \[=\langle\mathbf{w}_{y_{i},r}^{(\widetilde{t}-1,0)},\mathbf{\xi}_{i} \rangle-\frac{\eta}{Bm}\cdot\sum_{b^{\prime}=0}^{H-1}\sum_{i^{\prime}\in \mathcal{I}_{\widetilde{t}-1,b^{\prime}}}\ell_{i^{\prime}}^{(\widetilde{t}-1,b ^{\prime})}\cdot\sigma^{\prime}(\langle\mathbf{w}_{y_{i},r}^{(\widetilde{t}-1, b^{\prime})},\widehat{y}_{i^{\prime}}\mathbf{\mu}\rangle)\cdot\langle y_{i^{\prime}}\mathbf{ \mu},\mathbf{\xi}_{i}\rangle y_{i^{\prime}}\] \[\qquad-\frac{\eta(P-1)}{Bm}\cdot\ell_{i}^{(\widetilde{t}-1,b_{i} ^{(\widetilde{t}-1)})}\cdot\sigma^{\prime}(\langle\mathbf{w}_{y_{i},r}^{( \widetilde{t}-1,\widetilde{b})},\mathbf{\xi}_{i}\rangle)\cdot\|\mathbf{\xi}_{i}\|_{2}^ {2}\] \[\qquad-\frac{\eta(P-1)}{Bm}\cdot\sum_{b^{\prime}=0}^{H-1}\sum_{i^{ \prime}\in\mathcal{I}_{\widetilde{t}-1,b^{\prime}}}\ell_{i^{\prime}}^{( \widetilde{t},b^{\prime})}\cdot\sigma^{\prime}(\langle\mathbf{w}_{y_{i},r}^{( \widetilde{t},b^{\prime})},\mathbf{\xi}_{i^{\prime}}\rangle)\cdot\langle\mathbf{\xi}_{i^ {\prime}},\mathbf{\xi}_{i}\rangle\,\mathds{1}(i^{\prime}\neq i)\] \[=\langle\mathbf{w}_{y_{i},r}^{(\widetilde{t},0)},\mathbf{\xi}_{i} \rangle-\frac{\eta(P-1)}{Bm}\cdot\underbrace{\ell_{i}^{(\widetilde{t}-1,b_{i} ^{(\widetilde{t}-1)})}\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}}_{\text{I}_{4}}\] \[\qquad-\frac{\eta(P-1)}{Bm}\cdot\underbrace{\sum_{b^{\prime}=0}^{H -1}\sum_{i^{\prime}\in\mathcal{I}_{\widetilde{t}-1,b^{\prime}}}\ell_{i^{\prime}}^{( \widetilde{t}-1,b^{\prime})}\cdot\sigma^{\prime}(\langle\mathbf{w}_{y_{i},r}^{( \widetilde{t}-1,b^{\prime})},\mathbf{\xi}_{i^{\prime}}\rangle)\cdot\langle\mathbf{\xi}_{i^ {\prime}},\mathbf{\xi}_{i}\rangle\,\mathds{1}(i^{\prime}\neq i)}_{\text{I}_{5}}\]\[-\frac{\eta}{Bm}\cdot\underbrace{\sum_{\vec{b}=0}^{H-1}\sum_{i^{\prime}\in\mathcal{I} _{t-1,\nu}}\ell_{i^{\prime}}^{(\tilde{t}-1,b^{\prime})}\cdot\sigma^{\prime}( \langle\mathbf{w}_{y_{i,r}}^{(\tilde{t}-1,b^{\prime})},y_{i^{\prime}}\mathbf{\mu} \rangle)\cdot\langle y_{i^{\prime}}\mathbf{\mu},\mathbf{\xi}_{i}\rangle y_{i^{\prime}}} _{I_{6}},\]

for any \(r\in S_{i}^{(\tilde{t}-1,0)}\), where the last equality is by \(\langle\mathbf{w}_{y_{i,r}}^{(\tilde{t}-1,b_{i}^{(\tilde{t}-1)})},\mathbf{\xi}_{i} \rangle>0\). Then we respectively estimate \(\mathrm{I}_{4},\mathrm{I}_{5},\mathrm{I}_{6}\). For \(\mathrm{I}_{4}\), according to Lemma B.1, we have

\[-\mathrm{I}_{4}\geq|\ell_{i}^{(\tilde{t}-1,b_{i}^{(\tilde{t}-1)})}|\cdot \sigma_{p}^{2}d/2.\]

For \(\mathrm{I}_{5}\), we have following upper bound

\[|\mathrm{I}_{5}| \leq\sum_{b^{\prime}=0}^{H-1}\sum_{i^{\prime}\in\mathcal{I}_{ \tilde{t}-1,b^{\prime}}}|\ell_{i^{\prime}}^{(\tilde{t}-1,b^{\prime})}|\cdot \sigma^{\prime}(\langle\mathbf{w}_{y_{i,r}}^{(\tilde{t}-1,b^{\prime})},\mathbf{\xi }_{i^{\prime}}\rangle)\cdot|\langle\mathbf{\xi}_{i^{\prime}},\mathbf{\xi}_{i}\rangle| \operatorname{\mathds{1}}(i^{\prime}\neq i)\] \[\leq\sum_{b^{\prime}=0}^{H-1}\sum_{i^{\prime}\in\mathcal{I}_{ \tilde{t}-1,b^{\prime}}}|\ell_{i^{\prime}}^{(\tilde{t}-1,b^{\prime})}|\cdot| \langle\mathbf{\xi}_{i^{\prime}},\mathbf{\xi}_{i}\rangle|\operatorname{\mathds{1}}(i^{ \prime}\neq i)\] \[\leq\sum_{b^{\prime}=0}^{H-1}\sum_{i^{\prime}\in\mathcal{I}_{ \tilde{t}-1,b^{\prime}}}|\ell_{i^{\prime}}^{(\tilde{t}-1,b^{\prime})}|\cdot 2 \sigma_{p}^{2}\cdot\sqrt{d\log(6n^{2}/\delta)}\] \[\leq nC_{2}\big{|}\ell_{i}^{(\tilde{t}-1,b_{i}^{(\tilde{t}-1)})} \big{|}\cdot 2\sigma_{p}^{2}\cdot\sqrt{d\log(6n^{2}/\delta)},\]

where the first inequality is due to triangle inequality, the second inequality is due to \(\sigma^{\prime}(z)\in\{0,1\}\), the third inequality is due to Lemma B.1, the forth inequality is due to the third hypothesis at epoch \(\tilde{t}-1\).

For \(\mathrm{I}_{6}\), we have following upper bound

\[|\mathrm{I}_{6}| \leq\sum_{b^{\prime}=0}^{H-1}\sum_{i^{\prime}\in\mathcal{I}_{ \tilde{t}-1,b^{\prime}}}|\ell_{i^{\prime}}^{(\tilde{t}-1,b^{\prime})}|\cdot \sigma^{\prime}(\langle\mathbf{w}_{y_{i,r}}^{(\tilde{t}-1,b^{\prime})},y_{i^{ \prime}}\mathbf{\mu}\rangle)\cdot|\langle y_{i^{\prime}}\mathbf{\mu},\mathbf{\xi}_{i}\rangle|\] \[\leq\sum_{b^{\prime}=0}^{H-1}\sum_{i^{\prime}\in\mathcal{I}_{ \tilde{t}-1,b^{\prime}}}|\ell_{i^{\prime}}^{(\tilde{t}-1,b^{\prime})}||\langle y _{i^{\prime}}\mathbf{\mu},\mathbf{\xi}_{i}\rangle|\] \[\leq nC_{2}\big{|}\ell_{i}^{(\tilde{t}-1,b_{i}^{(\tilde{t}-1)})} \big{|}\cdot\|\mathbf{\mu}\|_{2}\sigma_{p}\sqrt{2\log(6n/\delta)},\]

where the first inequality is by triangle inequality; the second inequality is due to \(\sigma^{\prime}(z)\in\{0,1\}\); the third inequality is by Lemma B.1; the last inequality is due to the third hypothesis at epoch \(\tilde{t}-1\).

Since \(d\geq\max\{32C_{2}^{2}n^{2}\cdot\log(6n^{2}/\delta),4C_{2}n\|\mathbf{\mu}\|\sigma_ {p}^{-1}\sqrt{2\log(6n/\delta)}\}\), we have \(-(P-1)\mathrm{I}_{4}\geq\max\{(P-1)|\mathrm{I}_{5}|/2,|\mathrm{I}_{6}|/2\}\) and hence \(-(P-1)\mathrm{I}_{4}\geq(P-1)|\mathrm{I}_{5}|+|\mathrm{I}_{6}|\). It follows that

\[\langle\mathbf{w}_{y_{i,r}}^{(\tilde{t},0)},\mathbf{\xi}_{i}\rangle\geq\langle \mathbf{w}_{y_{i,r}}^{(\tilde{t}-1,0)},\mathbf{\xi}_{i}\rangle>\sigma_{0}\sigma_{p} \sqrt{d}/\sqrt{2},\]

for any \(r\in S_{i}^{(\tilde{t}-1,0)}\). Therefore, \(S_{i}^{(0,0)}\subseteq S_{i}^{(\tilde{t}-1,0)}\subseteq S_{i}^{(\tilde{t},0)}\). And it directly follows by Lemma B.3 that \(|S_{i}^{(\tilde{t},0)}|\geq 0.8m\Phi(-1),\,\forall i\in[n]\).

For the fifth hypothesis, similar to the proof of the fourth hypothesis, we also have

\[\langle\mathbf{w}_{y_{i},r}^{(\tilde{t},0)},\mathbf{\xi}_{i}\rangle =\langle\mathbf{w}_{y_{i},r}^{(\tilde{t}-1,0)},\mathbf{\xi}_{i} \rangle-\frac{\eta(P-1)}{Bm}\cdot\underbrace{\ell_{i}^{(\tilde{t}-1,b_{i}^{(t -1)})}\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}}_{\mathrm{I}_{4}}\] \[\qquad-\frac{\eta(P-1)}{Bm}\cdot\underbrace{\sum_{b^{\prime}=0}^ {H-1}\sum_{i^{\prime}\in\mathcal{I}_{\tilde{t},b^{\prime}}}\ell_{i^{\prime}}^{( \tilde{t}-1,b^{\prime})}\cdot\sigma^{\prime}(\langle\mathbf{w}_{y_{i,r}}^{( \tilde{t}-1,b^{\prime})},\mathbf{\xi}_{i^{\prime}}\rangle)\cdot\langle\mathbf{\xi}_{i^{ \prime}},\mathbf{\xi}_{i}\rangle\operatorname{\mathds{1}}(i^{\prime}\neq i)}_{ \mathrm{I}_{5}}\]\[-\frac{\eta}{Bm}\cdot\underbrace{\sum_{\mathbf{\nu}=0}^{H-1}\sum_{i^{\prime}\in\mathcal{I }_{z-1,\mathbf{\nu}}}\ell_{i^{\prime}}^{(\bar{\ell}-1,b^{\prime})}\cdot\sigma^{\prime }(\langle\mathbf{w}_{y_{i},r}^{(\bar{\ell}-1,b^{\prime})},y_{i^{\prime}}\mathbf{ \mu}\rangle)\cdot\langle y_{i^{\prime}}\mathbf{\mu},\mathbf{\xi}_{i}\rangle y_{i^{ \prime}}}_{\mathbf{l}_{6}},\]

for any \(i\in S_{j,r}^{(\bar{\ell}-1,0)}\), where the equality holds due to \(\langle\mathbf{w}_{j,r}^{(\bar{\ell}-1,b_{i}^{(\bar{\ell}-1)})},\mathbf{\xi}_{i} \rangle>0\) and \(y_{i}=j\). By applying the same technique used in the proof of the fourth hypothesis, it follows that

\[\langle\mathbf{w}_{j,r}^{(\bar{\ell},0)},\mathbf{\xi}_{i}\rangle\geq\langle \mathbf{w}_{j,r}^{(\bar{\ell}-1,0)},\mathbf{\xi}_{i}\rangle>0,\]

for any \(i\in S_{j,r}^{(\bar{\ell}-1,0)}\). Thus, we have \(S_{j,r}^{(0,0)}\subseteq S_{j,r}^{(\bar{\ell}-1,0)}\subseteq S_{j,r}^{(\bar{ \ell},0)}\). And it directly follows by Lemma B.4 that \(|S_{j,r}^{(\bar{\ell},0)}|\geq n\Phi(-1)/4\).

Proof of Proposition c.2.: Our proof is based on induction. The results are obvious at iteration \((0,0)\) as all the coefficients are zero. Suppose that the results in Proposition C.2 hold for all iterations \((0,0)\leq(t,b)<(\widetilde{t},\widetilde{b})\). We aim to prove that they also hold for iteration \((\widetilde{t},\widetilde{b})\).

Firstly, We prove that (28) exists at iteration \((\widetilde{t},\widetilde{b})\), i.e., \(\underline{\rho}_{j,r,i}^{(\widetilde{t},\widetilde{b})}\geq-\beta-10\sqrt{ \log(6n^{2}/\delta)/d}\cdot n\alpha\) for any \(r\in[m]\), \(j\in\{\pm 1\}\) and \(i\in[n]\). Notice that \(\underline{\rho}_{j,r,i}^{(\widetilde{t},\widetilde{b})}=0\) for \(j=y_{i}\), therefore we only need to consider the case that \(j\neq y_{i}\). We also only need to consider the case of \(\widetilde{b}=b_{i}^{(\widetilde{t})}+1\) since \(\underline{\rho}_{j,r,i}^{(\widetilde{t},\widetilde{b})}\) doesn't change in other cases according to (21).

When \(\underline{\rho}_{j,r,t}^{(\widetilde{t},b_{i}^{(\widetilde{t})})}<-0.5\beta- 5\sqrt{\log(6n^{2}/\delta)/d}\cdot n\alpha\), by (32) in Lemma C.3 we have that

\[(P-1)\langle\mathbf{w}_{j,r}^{(\widetilde{t},b_{i}^{(\widetilde{t})})},\mathbf{ \xi}_{i}\rangle\leq\underline{\rho}_{j,r,i}^{(\widetilde{t},b_{i}^{(\widetilde {t})})}+(P-1)\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle+5\sqrt{\frac{ \log(6n^{2}/\delta)}{d}}n\alpha<0,\]

and thus

\[\underline{\rho}_{j,r,i}^{(\widetilde{t},\widetilde{b})} =\underline{\rho}_{j,r,i}^{(\widetilde{t},b_{i}^{(\widetilde{t}) })}+\frac{\eta(P-1)^{2}}{Bm}\cdot\ell_{i}^{(\widetilde{t},b_{i}^{(\widetilde {t})})}\cdot\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(\widetilde{t},b_{i}^{( \widetilde{t})})},\mathbf{\xi}_{i}\rangle)\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}\] \[=\underline{\rho}_{j,r,i}^{(\widetilde{t},b_{i}^{(\widetilde{t}) })}\geq-\beta-10\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha,\]

where the last inequality is by induction hypothesis.

When \(\underline{\rho}_{j,r,i}^{(\widetilde{t},b_{i}^{(\widetilde{t})})}\geq-0.5\beta -5\sqrt{\log(6n^{2}/\delta)/d}\cdot n\alpha\), we have

\[\underline{\rho}_{j,r,i}^{(\widetilde{t},\widetilde{b})} =\underline{\rho}_{j,r,i}^{(t,b_{i}^{(\widetilde{t})})}+\frac{ \eta(P-1)^{2}}{Bm}\cdot\ell_{i}^{(t,b_{i}^{(\widetilde{t})})}\cdot\sigma^{ \prime}(\langle\mathbf{w}_{j,r}^{(t,b_{i}^{(\widetilde{t})})},\mathbf{\xi}_{i} \rangle)\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}\] \[\geq-0.5\beta-5\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha-\frac {\eta(P-1)^{2}\cdot 3\sigma_{p}^{2}d}{2Bm}\] \[\geq-0.5\beta-10\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha\] \[\geq-\beta-10\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha,\]

where the first inequality is by \(\ell_{i}^{\prime(t,b_{i}^{(\widetilde{t})})}\in(-1,0)\) and \(\|\mathbf{\xi}_{i}\|_{2}^{2}\leq(3/2)\sigma_{p}^{2}d\) by Lemma B.1; the second inequality is due to \(5\sqrt{\log(6n^{2}/\delta)/d}\cdot n\alpha\geq 3\eta\sigma_{p}^{2}d/(2Bm)\) by Condition 3.1.

Next we prove (27) holds for \((\widetilde{t},\widetilde{b})\). We only need to consider the case of \(j=y_{i}\). Consider

\[|\ell_{i}^{(\widetilde{t},\widetilde{b})}| =\frac{1}{1+\exp\{y_{i}\cdot[F_{+1}(\mathbf{W}_{+1}^{(\widetilde {t},\widetilde{b})},\mathbf{x}_{i})-F_{-1}(\mathbf{W}_{-1}^{(\widetilde{t}, \widetilde{b})},\mathbf{x}_{i})]\}} \tag{40}\] \[\leq\exp(-y_{i}\cdot[F_{+1}(\mathbf{W}_{+1}^{(\widetilde{t}, \widetilde{b})},\mathbf{x}_{i})-F_{-1}(\mathbf{W}_{-1}^{(\widetilde{t}, \widetilde{b})},\mathbf{x}_{i})])\] \[\leq\exp(-F_{y_{i}}(\mathbf{W}_{y_{i}}^{(\widetilde{t},\widetilde {b})},\mathbf{x}_{i})+0.5),\]where the last inequality is by \(F_{j}(\mathbf{W}_{j}^{(t,b)},\mathbf{x}_{i})\leq 0.5\) for \(j\neq y_{i}\) according to Lemma C.4. Now recall the iterative update rule of \(\overline{\rho}_{j,r,i}^{(t,b)}\):

\[\overline{\rho}_{j,r,i}^{(t,b+1)}=\overline{\rho}_{j,r,i}^{(t,b)}-\frac{\eta(P -1)^{2}}{Bm}\cdot\ell_{i}^{\prime(t,b)}\cdot\sigma^{\prime}(\langle\mathbf{w}_{ j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle)\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}\cdot\mathds{1}(i \in\mathcal{I}_{t,b}).\]

Let \((t_{j,r,i},b_{j,r,i})\) be the last time before \((\widetilde{t},\widetilde{b})\) that \(\overline{\rho}_{j,r,i}^{(t_{j,r,i},b_{j,r,i})}\leq 0.5\alpha\). Then by iterating the update rule from \((t_{j,r,i},b_{j,r,i})\) to \((\widetilde{t},\widetilde{b})\), we get

\[\begin{split}&\overline{\rho}_{j,r,i}^{(\widetilde{t},\widetilde{ b})}\\ &=\overline{\rho}_{j,r,i}^{(t_{j,r,i},b_{j,r,i})}-\underbrace{ \eta(P-1)^{2}}_{Bm}\cdot\ell_{i}^{\prime(t_{j,r,i},b_{j,r,i})}\cdot\mathds{1}( \langle\mathbf{w}_{j,r}^{(t_{j,r,i},b_{j,r,i})},\mathbf{\xi}_{i}\rangle\geq 0)\cdot \mathds{1}(i\in\mathcal{I}_{t,b})\|\mathbf{\xi}_{i}\|_{2}^{2}\\ &\qquad-\underbrace{\sum_{(t_{j,r,i},b_{j,r,i})<(t,b)<(\widetilde {t},\widetilde{b})}\frac{\eta(P-1)^{2}}{Bm}\cdot\ell_{i}^{\prime(t,b)}\cdot \mathds{1}(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle\geq 0)\cdot \mathds{1}(i\in\mathcal{I}_{t,b})\|\mathbf{\xi}_{i}\|_{2}^{2}}_{1s}.\end{split} \tag{41}\]

We first bound \(\mathds{I}_{7}\) as follows:

\[|\mathds{I}_{7}|\leq(\eta(P-1)^{2}/Bm)\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}\leq(\eta(P -1)^{2}/Bm)\cdot 3\sigma_{p}^{2}d/2\leq 1\leq 0.25\alpha,\]

where the first inequality is by \(\ell_{i}^{\prime(t_{j,r,i},b_{j,r,i})}\in(-1,0)\); the second inequality is by Lemma B.1; the third inequality is by Condition 3.1; the last inequality is by our choice of \(\alpha=4\log(T^{*})\) and \(T^{*}\geq e\).

Second, we bound \(\mathds{I}_{8}\). For \((t_{j,r,i},b_{j,r,i})<(t,b)<(\widetilde{t},\widetilde{b})\) and \(y_{i}=j\), we can lower bound the inner product \(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle\) as follows

\[\begin{split}\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle& \geq\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle+\frac{1}{P -1}\overline{\rho}_{j,r,i}^{(t,b)}-\frac{5}{P-1}\sqrt{\frac{\log(6n^{2}/ \delta)}{d}}n\alpha\\ &\geq-\frac{0.5}{P-1}\beta+\frac{0.5}{P-1}\alpha-\frac{5}{P-1} \sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha\\ &\geq\frac{0.25}{P-1}\alpha,\end{split} \tag{42}\]

where the first inequality is by (31) in Lemma C.3; the second inequality is by \(\overline{\rho}_{j,r,i}^{(t,b)}>0.5\alpha\) and \(\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle\geq-0.5\beta/(P-1)\) due to the definition of \(t_{j,r,i}\) and \(\beta\); the last inequality is by \(\beta\leq 1/8\leq 0.1\alpha\) and \(5\sqrt{\log(6n^{2}/\delta)/d}\cdot n\alpha\leq 0.2\alpha\) by Condition 3.1.

Thus, plugging the lower bounds of \(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle\) into \(\mathds{I}_{8}\) gives

\[\begin{split}|\mathds{I}_{8}|&\leq\sum_{(t_{j,r,i},b_{j,r,i})<(t,b)}\frac{\eta(P-1)^{2}}{Bm}\cdot\exp\Big{(}-\frac{1}{m}\sum_{r= 1}^{m}(P-1)\sigma(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle)+0.5 \Big{)}\\ &\qquad\qquad\qquad\qquad\qquad\qquad\qquad\cdot\mathds{1}( \langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle\geq 0)\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}\\ &\leq\frac{2\eta\Gamma^{*}n(P-1)^{2}}{Bm}\cdot\exp(-0.25\alpha) \exp(0.5)\cdot\frac{3\sigma_{p}^{2}d}{2}\\ &\leq\frac{2\eta\Gamma^{*}n(P-1)^{2}}{Bm}\cdot\exp(-\log(T^{*})) \exp(0.5)\cdot\frac{3\sigma_{p}^{2}d}{2}\\ &=\frac{2\eta n(P-1)^{2}}{Bm}\cdot\frac{3\sigma_{p}^{2}d}{2}\exp( 0.5)\leq 1\leq 0.25\alpha,\end{split}\]

where the first inequality is by (40); the second inequality is by (42); the third inequality is by \(\alpha=4\log(T^{*})\); the fourth inequality is by Condition 3.1; the last inequality is by \(\log(T^{*})\geq 1\) and \(\alpha=4\log(T^{*})\). Plugging the bound of \(\mathds{I}_{7},\mathds{I}_{8}\) into (41) completes the proof for \(\overline{\rho}\).

For the upper bound of (29), we prove a augmented hypothesis that there exists a \(i^{*}\in[n]\) with \(y_{i^{*}}=j\) such that for \(1\leq t\leq T^{*}\) we have that \(\gamma_{j,r}^{(t,0)}/\overline{\rho}_{j,r,i^{*}}\leq C^{\prime}\widehat{\gamma}\). Recall the iterative update rule of \(\gamma_{j,r}^{(t,b)}\) and \(\overline{\rho}_{j,r,i}^{(t,b)}\), we have

\[\overline{\rho}_{j,r,i}^{(t,b+1)} =\overline{\rho}_{j,r,i}^{(t,b)}-\frac{\eta(P-1)^{2}}{Bm}\cdot \ell_{i}^{\prime(t,b)}\cdot\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)}, \boldsymbol{\xi}_{i}\rangle)\cdot\|\boldsymbol{\xi}_{i}\|_{2}^{2}\cdot\mathds{1 }(y_{i}=j)\,\mathds{1}(i\in\mathcal{I}_{t,b}),\] \[\gamma_{j,r}^{(t,b+1)} =\gamma_{j,r}^{(t,b)}-\frac{\eta}{Bm}\cdot\bigg{[}\sum_{i\in \mathcal{I}_{t,b}\cap S_{+}}\ell_{i}^{\prime(t,b)}\sigma^{\prime}(\langle \mathbf{w}_{j,r}^{(t,b)},y_{i}\cdot\boldsymbol{\mu}\rangle)\] \[\qquad\qquad\qquad\qquad\qquad-\sum_{i\in\mathcal{I}_{t,b}\cap S_ {-}}\ell_{i}^{\prime(t,b)}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},y_{ i}\cdot\boldsymbol{\mu}\rangle)\bigg{]}\cdot\|\boldsymbol{\mu}\|_{2}^{2}.\]

According to the fifth statement of Lemma C.8, for any \(i^{*}\in S_{j,r}^{(0,0)}\) it holds that \(j=y_{i^{*}}\) and \(\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\xi}_{i^{*}}\rangle\geq 0\) for any \((t,b)\leq(\widetilde{t},\widetilde{b})\). Thus, we have

\[\overline{\rho}_{j,r,i^{*}}^{(\widetilde{t},0)}=\overline{\rho}_{j,r,i^{*}}^{( \widetilde{t}-1,0)}-\frac{\eta(P-1)^{2}}{Bm}\cdot\ell_{i^{*}}^{(\widetilde{t} -1,b_{i^{*}}^{(\widetilde{t}-1)})}\cdot\|\boldsymbol{\xi}_{i^{*}}\|_{2}^{2} \geq\overline{\rho}_{j,r,i^{*}}^{(\widetilde{t}-1,0)}-\frac{\eta(P-1)^{2}}{Bm }\cdot\ell_{i^{*}}^{(\widetilde{t}-1,b_{i^{*}}^{(\widetilde{t}-1)})}\cdot \sigma_{p}^{2}d/2.\]

For the update rule of \(\gamma_{j,r}^{(t,b)}\), according to Lemma C.8, we have

\[\sum_{b<H}\Big{|}\sum_{i\in\mathcal{I}_{t,b}\cap S_{+}}\ell_{i}^{ \prime(t,b)}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},y_{i}\cdot \boldsymbol{\mu}\rangle) -\sum_{i\in\mathcal{I}_{t,b}\cap S_{-}}\ell_{i}^{\prime(t,b)} \sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},y_{i}\cdot\boldsymbol{\mu} \rangle)\Big{|}\] \[\leq C_{2}n\Big{|}\ell_{i^{*}}^{(\widetilde{T}-1,b_{i^{*}}^{( \widetilde{T}-1)})}\Big{|}.\]

Then, we have

\[\frac{\gamma_{j,r}^{(\widetilde{t},0)}}{\overline{\rho}_{j,r,i^{*}}^{( \widetilde{t},0)}} \leq\max\left\{\frac{\gamma_{j,r}^{(\widetilde{t}-1,0)}}{\overline {\rho}_{j,r,i^{*}}^{(\widetilde{t}-1,0)}},\frac{C_{2}n\ell_{i^{*}}^{\prime( \widetilde{t}-1,b_{i^{*}}^{(\widetilde{t}-1)})}\|\boldsymbol{\mu}\|_{2}^{2}}{ (P-1)^{2}\cdot\ell_{i^{*}}^{\prime(\widetilde{t}-1,b_{i^{*}}^{(\widetilde{t}- 1)})}\cdot\sigma_{p}^{2}d/2}\right\}\] \[=\max\left\{\frac{\gamma_{j,r}^{(\widetilde{t}-1,0)}}{\overline{ \rho}_{j,r,i^{*}}^{(\widetilde{t}-1,0)}},\frac{2C_{2}n\|\boldsymbol{\mu}\|_{2} ^{2}}{(P-1)^{2}\sigma_{p}^{2}d}\right\} \tag{43}\] \[\leq\frac{2C_{2}n\|\boldsymbol{\mu}\|_{2}^{2}}{(P-1)^{2}\sigma_{ p}^{2}d},\]

where the last inequality is by \(\gamma_{j,r}^{(\widetilde{t}-1,0)}/\overline{\rho}_{j,r,i^{*}}^{(\widetilde{t}-1,0)}\leq 2 C_{2}\widehat{\gamma}=2C_{2}n\|\boldsymbol{\mu}\|_{2}^{2}/(P-1)^{2}\sigma_{p}^ {2}d\). Therefore,

\[\frac{\gamma_{j,r}^{(\widetilde{t},0)}}{\overline{\rho}_{j,r,i^{*}}^{(\widetilde {t},0)}}\leq 2C_{2}\widehat{\gamma}.\]

For iterations other than the starting of en epoch, we have the following upper bound:

\[\frac{\gamma_{j,r}^{(\widetilde{t},b)}}{\overline{\rho}_{j,r,i^{*}}^{( \widetilde{t},b)}}\leq\frac{2\gamma_{j,r}^{(\widetilde{t},0)}}{\overline{\rho }_{j,r,i^{*}}^{(\widetilde{t},0)}}\leq 4C_{2}\widehat{\gamma}.\]

Thus, by taking \(C^{\prime}=4C_{2}\), we have \(\gamma_{j,r}^{(\widetilde{t},b)}/\overline{\rho}_{j,r,i^{*}}^{(\widetilde{t}, b)}\leq C^{\prime}\widehat{\gamma}\).

On the other hand, when \((t,b)<(\frac{\log(2T^{*}/\delta)}{2c_{3}^{3}},0)\), we have

\[\gamma_{j,r}^{(t,b)}\geq-\frac{\log(2T^{*}/\delta)}{2c_{3}^{3}}\cdot\frac{\eta} {Bm}\cdot n\cdot\|\boldsymbol{\mu}\|_{2}^{2}\geq-\frac{1}{12},\]

where the first inequality is due to update rule of \(\gamma_{j,r}^{t,b}\), and the second inequality is due to Condition 3.1.

When \((t,b)\geq\big{(}\frac{\log(2T^{7}/\delta)}{2c_{3}^{3}},0\big{)}\), According to Lemma B.6, we have

\[\gamma_{j,r}^{(t,b)} \geq\sum_{(t^{\prime},b^{\prime})<(t,b)}\frac{\eta}{Bm}\big{[}\min _{i,b^{\prime}}\ell_{i}^{\prime(t^{\prime},b^{\prime})}\min\{|\mathcal{I}_{t^{ \prime},b^{\prime}}\cap S_{+}\cap S_{-1}|,|\mathcal{I}_{t^{\prime},b^{\prime}} \cap S_{+}\cap S_{1}|\}\] \[\qquad\qquad\qquad\qquad\qquad-\max_{i,b^{\prime}}\ell_{i}^{(t^{ \prime},b^{\prime})}|\mathcal{I}_{t^{\prime},b^{\prime}}\cap S_{-}|\big{]} \cdot\|\mathbf{\mu}\|_{2}^{2}\] \[\geq\frac{\eta}{Bm}\big{(}\sum_{t^{\prime}=0}^{t-1}(c_{3}c_{4}H \frac{B}{4}\min_{i,b^{\prime}}\ell_{i}^{\prime(t^{\prime},b^{\prime})}-nq\max_ {i,b^{\prime}}\ell_{i}^{\prime(t^{\prime},b^{\prime})})-nq\max_{i,b^{\prime}} \ell_{i}^{(t,b^{\prime})}\big{)}\big{)}\|\mathbf{\mu}\|_{2}^{2}\] \[\geq 0,\]

where the first inequality is due to the update rule of \(\gamma_{j,r}^{(t,b)}\), the second inequality is due to Lemma B.6, and the third inequality is due to Condition 3.1.

### Decoupling with a Two-stage Analysis

#### c.2.1 First Stage

**Lemma C.9**.: _There exist_

\[T_{1}=C_{3}\eta^{-1}Bm(P-1)^{-2}\sigma_{p}^{-2}d^{-1},T_{2}=C_{4}\eta^{-1}Bm(P -1)^{-2}\sigma_{p}^{-2}d^{-1},\]

_where \(C_{3}=\Theta(1)\) is a large constant and \(C_{4}=\Theta(1)\) is a small constant, such that_

* \(\overline{\rho}_{j,r^{*},i}^{(T_{1},0)}\geq 2\) _for any_ \(r^{*}\in S_{i}^{(0,0)}=\{r\in[m]:\langle\mathbf{w}_{y_{i},r}^{(0)},\mathbf{\xi}_{i }\rangle>0\}\)_,_ \(j\in\{\pm 1\}\) _and_ \(i\in[n]\) _with_ \(y_{i}=j\)_._
* \(\max_{j,r}\gamma_{j,r}^{(t,b)}=O(\widehat{\gamma})\) _for all_ \((t,b)\leq(T_{1},0)\)_._
* \(\max_{j,r,i}|\varrho_{j,r,i}^{(t,b)}|=\max\{\beta,O\big{(}n\sqrt{\log(n/\delta )}\log(T^{*})/\sqrt{d}\big{)}\}\) _for all_ \((t,b)\leq(T_{1},0)\)_._
* \(\min_{j,r}\gamma_{j,r}^{(t,0)}=\Omega(\widehat{\gamma})\) _for all_ \(t\geq T_{2}\)_._
* \(\max_{j,r}\overline{\rho}_{j,r,i}^{(T_{1},0)}=O(1)\) _for all_ \(i\in[n]\)_._

Proof of Lemma c.9.: By Proposition C.2, we have that \(\underline{\rho}_{j,r,i}^{(t,b)}\geq-\beta-10n\sqrt{\frac{\log(6n^{2}/\delta )}{d}}\alpha\) for all \(j\in\{\pm 1\}\), \(r\in[m]\), \(i\in[n]\) and \((0,0)\leq(t,b)\leq(T^{*},0)\). According to Lemma B.2, for \(\beta\) we have

\[\beta =2\max_{i,j,r}\{|\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\mu}\rangle|,(P-1)|\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle|\}\] \[\leq 2\max\{\sqrt{2\log(12m/\delta)}\cdot\sigma_{0}\|\mathbf{\mu}\|_{2 },2\sqrt{\log(12mn/\delta)}\cdot\sigma_{0}(P-1)\sigma_{p}\sqrt{d}\}\] \[=O\big{(}\sqrt{\log(mn/\delta)}\cdot\sigma_{0}(P-1)\sigma_{p} \sqrt{d}\big{)},\]

where the last equality is by the first condition of Condition 3.1. Since \(\underline{\rho}_{j,r,i}^{(t,b)}\leq 0\), we have that

\[\max_{j,r,i}|\underline{\rho}_{j,r,i}^{(t,b)}| =\max_{j,r,i}-\underline{\rho}_{j,r,i}^{(t,b)}\] \[\leq\beta+10\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha\] \[=\max\bigg{\{}\beta,O\big{(}\sqrt{\log(n/\delta)}\log(T^{*})\cdot n /\sqrt{d}\big{)}\bigg{\}}.\]

Next, for the growth of \(\gamma_{j,r}^{(t)}\), we have following upper bound

\[\gamma_{j,r}^{(t,b+1)} =\gamma_{j,r}^{(t,b)}-\frac{\eta}{Bm}\cdot\sum_{i\in\mathcal{I}_{ t,b}}\ell_{i}^{\prime(t,b)}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},y_{i} \cdot\mathbf{\mu}\rangle)\cdot\|\mathbf{\mu}\|_{2}^{2}\] \[\leq\gamma_{j,r}^{(t,b)}+\frac{\eta}{m}\cdot\|\mathbf{\mu}\|_{2}^{2},\]where the inequality is by \(|\ell^{\prime}|\leq 1\). Note that \(\gamma^{(0,0)}_{j,r}=0\) and recursively use the inequality \(tB+b\) times we have

\[\gamma^{(t,b)}_{j,r}\leq\frac{\eta(tH+b)}{m}\cdot\|\mathbf{\mu}\|_{2}^{2}. \tag{44}\]

Since \(n\cdot\mathrm{SNR}^{2}=n\|\mathbf{\mu}\|_{2}^{2}/\big{(}(P-1)^{2}\sigma_{p}^{2}d \big{)}=\widehat{\gamma}\), we have

\[T_{1}=C_{3}\eta^{-1}Bm(P-1)^{-2}\sigma_{p}^{-2}d^{-1}=C_{3}\eta^{-1}m\|\mathbf{\mu }\|_{2}^{-2}\widehat{\gamma}B/n.\]

And it follows that

\[\gamma^{(t)}_{j,r}\leq\frac{\eta(tH+b)}{m}\cdot\|\mathbf{\mu}\|_{2}^{2}\leq\frac{ \eta nT_{1}}{mB}\cdot\|\mathbf{\mu}\|_{2}^{2}\leq C_{3}\widehat{\gamma},\]

for all \((0,0)\leq(t,b)\leq(T_{1},0)\).

For \(\overline{\rho}^{(t)}_{j,r,i}\), recall from (20) that

\[\overline{\rho}^{(t+1,0)}_{y_{i},r,i}=\overline{\rho}^{(t,0)}_{y_{i},r,i}- \frac{\eta(P-1)^{2}}{Bm}\cdot\ell^{(t,b^{(t)}_{i})}_{i}\cdot\sigma^{\prime}( \langle\mathbf{w}^{(t,b^{(t)}_{i})}_{y_{i},r},\mathbf{\xi}_{i}\rangle)\cdot\|\mathbf{ \xi}_{i}\|_{2}^{2}.\]

According to Lemma C.8, for any \(r^{*}\in S^{(0,0)}_{i}=\{r\in[m]:\langle\mathbf{w}^{(0)}_{y_{i},r},\mathbf{\xi}_{i }\rangle>\sigma_{0}\sigma_{p}\sqrt{d}/\sqrt{2}\}\), we have \(\langle\mathbf{w}^{(t,b)}_{y_{i},r^{*}},\mathbf{\xi}_{i}\rangle>0\) for all \((0,0)\leq(t,b)\leq(T^{*},0)\) and hence

\[\overline{\rho}^{(t+1,0)}_{j,r^{*},i}=\overline{\rho}^{(t,0)}_{j,r^{*},i}- \frac{\eta(P-1)^{2}}{Bm}\cdot\ell^{(t,b^{(t)}_{i})}_{i}\|\mathbf{\xi}_{i}\|_{2}^{2}.\]

For each \(i\), we denote by \(T_{1}^{(i)}\) the last time in the period \([0,T_{1}]\) satisfying that \(\overline{\rho}^{(t,0)}_{y_{i},r^{*},i}\leq 2\). Then for \((0,0)\leq(t,b)<(T_{1}^{(i)},0)\), \(\max_{j,r}\{|\overline{\rho}^{(t,b)}_{j,r,i}|,|\varrho^{(t,b)}_{j,r,i}|\}=O(1)\) and \(\max_{j,r}\gamma^{(t,b)}_{j,r}=O(1)\). Therefore, we know that \(F_{-1}(\mathbf{W}^{(t,b)},\mathbf{x}_{i}),F_{+1}(\mathbf{W}^{(t,b)},\mathbf{x }_{i})=O(1)\). Thus there exists a positive constant \(C\) such that \(-\ell^{\prime(t,b)}_{i}\geq C\geq C_{2}\) for \(0\leq t\leq T_{1}^{(i)}\).

Then we have

\[\overline{\rho}^{(t,0)}_{y_{i},r^{*},i}\geq\frac{C\eta(P-1)^{2}\sigma_{p}^{2} dt}{2Bm}.\]

Therefore, \(\overline{\rho}^{(t,0)}_{y_{i},r^{*},i}\) will reach 2 within

\[T_{1}=C_{3}\eta^{-1}Bm(P-1)^{2}\sigma_{p}^{-2}d^{-1}\]

iterations for any \(r^{*}\in S^{(0,0)}_{i}\), where \(C_{3}\) can be taken as \(4/C\).

Next, we will discuss the lower bound of the growth of \(\gamma^{(t,b)}_{j,r}\). For \(\overline{\rho}^{(t,b)}_{j,r,i}\), we have

\[\overline{\rho}^{(t,b+1)}_{j,r,i} =\overline{\rho}^{(t,b)}_{j,r,i}-\frac{\eta(P-1)^{2}}{Bm}\cdot\ell ^{\prime(t,b)}_{i}\cdot\sigma^{\prime}(\langle\mathbf{w}^{(t,b)}_{j,r},\mathbf{ \xi}_{i}\rangle)\cdot\|\mathbf{\xi}_{i}\|_{2}^{2}\cdot\mathds{1}(y_{i}=j)\, \mathds{1}(i\in\mathcal{I}_{t,b})\] \[\leq\overline{\rho}^{(t,b)}_{j,r,i}+\frac{3\eta(P-1)^{2}\sigma_{p }^{2}d}{2Bm}.\]

According to (44) and \(\overline{\rho}^{(0,0)}_{j,r,i}=0\), it follows that

\[\overline{\rho}^{(t,b)}_{j,r,i}\leq\frac{3\eta(P-1)^{2}\sigma_{p}^{2}d(tH+b)} {2Bm},\gamma^{(t,b)}_{j,r}\leq\frac{\eta(tH+b)}{m}\cdot\|\mathbf{\mu}\|_{2}^{2}. \tag{45}\]

Therefore, \(\max_{j,r,i}\overline{\rho}^{(t,b)}_{j,r,i}\) will be smaller than \(1\) and \(\gamma^{(t,b)}_{j,r}\) smaller than \(\Theta(n\|\mathbf{\mu}\|_{2}^{2}/(P-1)^{2}\sigma_{p}^{2}d)=\Theta(n\cdot\mathrm{ SNR}^{2})=\Theta(\widehat{\gamma})=O(1)\) within

\[T_{2}=C_{4}\eta^{-1}Bm(P-1)^{-2}\sigma_{p}^{-2}d^{-1}\]

iterations, where \(C_{4}\) can be taken as \(2/3\). Therefore, we know that \(F_{-1}(\mathbf{W}^{(t,b)},\mathbf{x}_{i}),F_{+1}(\mathbf{W}^{(t,b)},\mathbf{x}_ {i})=O(1)\) in \((0,0)\leq(t,b)\leq(T_{2},0)\). Thus, there exists a positive constant \(C\) such that \(-\ell^{\prime(t,b)}_{i}\geq C\) for \(0\leq t\leq T_{2}\).

Recall that we denote \(\{i\in[n]|y_{i}=y\}\) as \(S_{y}\), and we have the update rule

\[\gamma_{j,r}^{(t,b+1)}=\gamma_{j,r}^{(t,b)}-\frac{\eta}{Bm}\cdot \bigg{[}\sum_{i\in\mathcal{I}_{k,b}\cap S_{+}}\ell_{i}^{\prime(t,b)} \sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\widehat{y}_{i}\cdot\mathbf{\mu} \rangle)\] \[\qquad\qquad\qquad\qquad\qquad-\sum_{i\in\mathcal{I}_{k,b}\cap S_{- }}\ell_{i}^{\prime(t,b)}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)}, \widehat{y}_{i}\cdot\mathbf{\mu}\rangle)\bigg{]}\cdot\|\mathbf{\mu}\|_{2}^{2}.\]

For the growth of \(\gamma_{j,r}^{(t,b)}\), if \(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\mu}\rangle\geq 0\), we have

\[\gamma_{j,r}^{(t,b+1)} =\gamma_{j,r}^{(t,b)}-\frac{\eta}{Bm}\cdot\bigg{[}\sum_{i\in \mathcal{I}_{k,b}\cap S_{+}\cap S_{1}}\ell_{i}^{\prime(t)}-\sum_{i\in \mathcal{I}_{k,b}\cap S_{-}\cap S_{1}}\ell_{i}^{\prime(t)}\bigg{]}\|\mathbf{\mu}\|_ {2}^{2}\] \[\geq\gamma_{j,r}^{(t,b)}+\frac{\eta}{Bm}\cdot\big{[}C|\mathcal{I} _{t,b}\cap S_{+}\cap S_{1}|-|\mathcal{I}_{t,b}\cap S_{-}\cap S_{-1}|\big{]} \cdot\|\mathbf{\mu}\|_{2}^{2}.\]

Similarly, if \(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\mu}\rangle<0\),

\[\gamma_{j,r}^{(t,b+1)}\geq\gamma_{j,r}^{(t,b)}+\frac{\eta}{Bm}\cdot\big{[}C| \mathcal{I}_{t,b}\cap S_{+}\cap S_{-1}|-|\mathcal{I}_{t,b}\cap S_{-}\cap S_{1 }|\big{]}\cdot\|\mathbf{\mu}\|_{2}^{2}.\]

Therefore, for \(t\in[T_{2},T_{1}]\), we have

\[\gamma_{j,r}^{(t,0)} \geq\sum_{(t^{\prime},b^{\prime})<(t,0)}\frac{\eta}{Bm}\big{[}C \min\{|\mathcal{I}_{t^{\prime},b^{\prime}}\cap S_{+}\cap S_{-1}|,|\mathcal{I} _{t^{\prime},b^{\prime}}\cap S_{+}\cap S_{1}|\}-|\mathcal{I}_{t,b}\cap S_{-}| \big{]}\cdot\|\mathbf{\mu}\|_{2}^{2}\] \[\geq\frac{\eta}{Bm}(c_{3}tc_{4}HC\frac{B}{4}-T_{1}nq)\|\mathbf{\mu}\| _{2}^{2}\] \[=\frac{\eta}{Bm}(c_{3}c_{4}TC\frac{n}{4}-T_{1}nq)\|\mathbf{\mu}\|_{2} ^{2}\] \[\geq\frac{\eta c_{3}c_{4}Ctn\|\mu\|_{2}^{2}}{8Bm} \tag{46}\] \[\geq\frac{c_{3}c_{4}CC_{4}n\|\mu\|_{2}^{2}}{(P-1)^{2}\sigma_{p}^ {2}d}=\Theta(n\cdot\mathrm{SNR}^{2})=\Theta(\widehat{\gamma}),\]

where the second inequality is due to Lemma B.6, the third inequality is due to \(q<\frac{C_{4}C_{4}c_{4}}{8C_{3}}\) in Condition 3.1.

And it follows directly from (45) that

\[\overline{\rho}_{j,r,i}^{(T_{1},0)}\leq\frac{3\eta(P-1)^{2}\sigma_{p}^{2}dT_{1 }H}{2Bm}=\frac{3C_{3}}{2},\overline{\rho}_{j,r,i}^{(T_{1},0)}=O(1),\]

which completes the proof.

#### c.2.2 Second Stage

By the signal-noise decomposition, at the end of the first stage, we have

\[\mathbf{w}_{j,r}^{(t,b)}=\mathbf{w}_{j,r}^{(0,0)}+j\gamma_{j,r}^{(t,b)}\|\mathbf{\mu }\|_{2}^{-2}\mathbf{\mu}+\frac{1}{P-1}\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(t,b )}\|\mathbf{\xi}_{i}\|_{2}^{-2}\mathbf{\xi}_{i}+\frac{1}{P-1}\sum_{i=1}^{n}\varrho_{j,r,i}^{(t,b)}\|\mathbf{\xi}_{i}\|_{2}^{-2}\mathbf{\xi}_{i}.\]

for \(j\in[\pm 1]\) and \(r\in[m]\). By the results we get in the first stage, we know that at the beginning of this stage, we have the following property holds:

* \(\overline{\rho}_{j,r^{*},i}^{(T_{1},0)}\geq 2\) for any \(r^{*}\in S_{i}^{(0,0)}=\{r\in[m]:\langle\mathbf{w}_{y_{i,r}}^{(0,0)},\mathbf{\xi}_ {i}\rangle>\sigma_{0}\sigma_{p}\sqrt{d}/\sqrt{2}\}\), \(j\in\{\pm 1\}\) and \(i\in[n]\) with \(y_{i}=j\).
* \(\max_{j,r,i}|\varrho_{j,r,i}^{(T_{1},0)}|=\max\{\beta,O\big{(}n\sqrt{\log(n/ \delta)}\log(T^{*})/\sqrt{d}\big{)}\}\).
* \(\gamma_{j,r}^{(T_{1},0)}=\Theta(\widehat{\gamma})\) for any \(j\in\{\pm 1\},r\in[m]\).

where \(\widehat{\gamma}=n\cdot\mathrm{SNR}^{2}\). Now we choose \(\mathbf{W}^{*}\) as follows

\[\mathbf{w}^{*}_{j,r}=\mathbf{w}^{(0,0)}_{j,r}+\frac{20\log(2/\epsilon)}{P-1} \Big{[}\sum_{i=1}^{n}\mathds{1}(j=y_{i})\cdot\frac{\mathbf{\xi}_{i}}{\|\mathbf{\xi}_{i} \|_{2}^{2}}\Big{]}.\]

**Lemma C.10**.: _Under the same conditions as Theorem 3.2, we have that \(\|\mathbf{W}^{(T_{1},0)}-\mathbf{W}^{*}\|_{F}\leq\widetilde{O}(m^{1/2}n^{1/2}(P -1)^{-1}\sigma_{p}^{-1}d^{-1/2}(1+\max\{\beta,n\sqrt{\log(n/\delta)}\log(T^{*} )/\sqrt{d}\}))\)._

Proof.: \[\|\mathbf{W}^{(T_{1},0)}-\mathbf{W}^{*}\|_{F}\] \[\leq\|\mathbf{W}^{(T_{1},0)}-\mathbf{W}^{(0,0)}\|_{F}+\|\mathbf{ W}^{*}-\mathbf{W}^{(0,0)}\|_{F}\] \[\leq O(\sqrt{m})\max_{j,r}\gamma_{j,r}^{(T_{1},0)}\|\mathbf{\mu}\|_{ 2}^{-1}+\frac{1}{P-1}O(\sqrt{m})\max_{j,r}\bigg{\|}\sum_{i=1}^{n}\overline{ \rho}_{j,r,i}^{(T_{1},0)}\cdot\frac{\mathbf{\xi}_{i}}{\|\mathbf{\xi}_{i}\|_{2}^{2}}+ \sum_{i=1}^{n}\underline{\rho}_{j,r,i}^{(T_{1},0)}\cdot\frac{\mathbf{\xi}_{i}}{\| \mathbf{\xi}_{i}\|_{2}^{2}}\bigg{\|}_{2}\] \[\qquad+O(m^{1/2}n^{1/2}\log(1/\epsilon)(P-1)^{-1}\sigma_{p}^{-1} d^{-1/2})\] \[=O(m^{1/2}\widehat{\gamma}\|\mathbf{\mu}\|_{2}^{-1})+\widetilde{O}(m ^{1/2}n^{1/2}(P-1)^{-1}\sigma_{p}^{-1}d^{-1/2}(1+\max\{\beta,n\sqrt{\log(n/ \delta)}\log(T^{*})/\sqrt{d}\}))\] \[\quad+O(m^{1/2}n^{1/2}\log(1/\epsilon)(P-1)^{-1}\sigma_{p}^{-1}d ^{-1/2})\] \[=O(m^{1/2}n\cdot\mathrm{SNR}\cdot(P-1)^{-1}\sigma_{p}^{-1}d^{-1/2 }(1+\max\{\beta,n\sqrt{\log(n/\delta)}\log(T^{*})/\sqrt{d}\}))\] \[\quad+\widetilde{O}(m^{1/2}n^{1/2}\log(1/\epsilon)(P-1)^{-1} \sigma_{p}^{-1}d^{-1/2})\] \[=\widetilde{O}(m^{1/2}n^{1/2}(P-1)^{-1}\sigma_{p}^{-1}d^{-1/2}(1+ \max\{\beta,n\sqrt{\log(n/\delta)}\log(T^{*})/\sqrt{d}\})),\]

where the first inequality is by triangle inequality, the second inequality and the first equality are by our decomposition of \(\mathbf{W}^{(T_{1},0)}\), \(\mathbf{W}^{*}\) and Lemma B.1; the second equality is by \(n\cdot\mathrm{SNR}^{2}=\Theta(\widehat{\gamma})\) and \(\mathrm{SNR}=\|\mathbf{\mu}\|/(P-1)\sigma_{p}d^{1/2}\); the third equality is by \(n^{1/2}\cdot\mathrm{SNR}=O(1)\).

**Lemma C.11**.: _Under the same conditions as Theorem 3.2, we have that_

\[y_{i}\langle\nabla f(\mathbf{W}^{(t,b)},\mathbf{x}_{i}),\mathbf{W}^{*}\rangle \geq\log(2/\epsilon)\]

_for all \((T_{1},0)\leq(t,b)\leq(T^{*},0)\)._

Proof of Lemma c.11.: Recall that \(f(\mathbf{W}^{(t,b)})=(1/m)\sum_{j,r}j\cdot[\sigma(\langle\mathbf{w}^{(t,b)}_{ j,r},\mathbf{\xi}_{i}\rangle)]\), thus we have

\[y_{i}\langle\nabla f(\mathbf{W}^{(t,b)},\mathbf{x}_{i}),\mathbf{W }^{*}\rangle\] \[=\frac{1}{m}\sum_{j,r}\sigma^{\prime}(\langle\mathbf{w}^{(t,b)}_{ j,r},\widehat{y}_{i}\mathbf{\mu}\rangle)\langle y_{i}\widehat{y}_{i}\mathbf{\mu},j \mathbf{w}^{*}_{j,r}\rangle+\frac{P-1}{m}\sum_{j,r}\sigma^{\prime}(\langle \mathbf{w}^{(t,b)}_{j,r},\mathbf{\xi}_{i}\rangle)\langle y_{i}\mathbf{\xi}_{i},j \mathbf{w}^{*}_{j,r}\rangle\] \[=\frac{1}{m}\sum_{j,r}\sum_{i^{\prime}=1}^{n}\sigma^{\prime}( \langle\mathbf{w}^{(t,b)}_{j,r},\mathbf{\xi}_{i}\rangle)20\log(2/\epsilon)\, \mathds{1}(j=y_{i^{\prime}})\cdot\frac{\langle\mathbf{\xi}_{i^{\prime}},\mathbf{\xi}_{ i}\rangle}{\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{2}}\] \[\qquad+\frac{1}{m}\sum_{j,r}\sum_{i^{\prime}=1}^{n}\sigma^{\prime }(\langle\mathbf{w}^{(t,b)}_{j,r},\widehat{y}_{i}\mathbf{\mu}\rangle)20\log(2/ \epsilon)\,\mathds{1}(j=y_{i^{\prime}})\cdot\frac{\langle\widehat{y}_{i}\mathbf{ \mu},\mathbf{\xi}_{i^{\prime}}\rangle}{\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{2}}\] \[\qquad+\frac{1}{m}\sum_{j,r}\sigma^{\prime}(\langle\mathbf{w}^{(t,b)}_{j,r},\widehat{y}_{i}\mathbf{\mu}\rangle)\langle y_{i}\widehat{y}_{i}\mathbf{\mu},j\mathbf{w}^{(0,0)}_{j,r}\rangle+\frac{P-1}{m}\sum_{j,r}\sigma^{\prime}( \langle\mathbf{w}^{(t,b)}_{j,r},\mathbf{\xi}_{i}\rangle)\langle y_{i}\mathbf{\xi}_{i}, j\mathbf{w}^{(0,0)}_{j,r}\rangle\] \[\geq\frac{1}{m}\sum_{j=y_{i},r}\sigma^{\prime}(\langle\mathbf{w} ^{(t,b)}_{j,r},\mathbf{\xi}_{i}\rangle)20\log(2/\epsilon)-\frac{1}{m}\sum_{j,r} \sum_{i^{\prime}\neq i}\sigma^{\prime}(\langle\mathbf{w}^{(t,b)}_{j,r},\mathbf{\xi} _{i}\rangle)20\log(2/\epsilon)\cdot\frac{|\langle\mathbf{\xi}_{i^{\prime}},\mathbf{\xi} _{i}\rangle|}{\|\mathbf{\xi}_{i^{\prime}}\|_{2}^{2}}\] \[\qquad-\frac{1}{m}\sum_{j,r}\sum_{i^{\prime}=1}^{n}\sigma^{\prime }(\langle\mathbf{w}^{(t,b)}_{j,r},\widehat{y}_{i}\mathbf{\mu}\rangle)20\log(2/ \epsilon)\cdot\frac{|\langle\widehat{y}_{i}\mathbf{\mu},\mathbf{\xi}_{i^{\prime}}\rangle|}{ \|\mathbf{\xi}_{i^{\prime}}\|_{2}^{2}}-\frac{1}{m}\sum_{j,r}\sigma^{\prime}( \langle\mathbf{w}^{(t,b)}_{j,r},\widehat{y}_{i}\mathbf{\mu}\rangle)\beta\]\[\geq\underbrace{\frac{1}{m}\sum_{j=y_{i},r}\sigma^{\prime}(\langle\mathbf{ w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle)20\log(2/\epsilon)}_{\mathbf{l}_{0}}- \underbrace{\frac{1}{m}\sum_{j,r}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b) },\mathbf{\xi}_{i}\rangle)20\log(2/\epsilon)O\big{(}n\sqrt{\log(n/\delta)}/\sqrt{d} \big{)}}_{\mathbf{l}_{1_{0}}}\] \[-\underbrace{\frac{1}{m}\sum_{j,r}\sigma^{\prime}(\langle\mathbf{ w}_{j,r}^{(t,b)},\widehat{y}_{i}\mathbf{\mu}\rangle)O\big{(}n\sqrt{\log(n/\delta)} \cdot\mathrm{SNR}\cdot d^{-1/2}\big{)}}_{\mathbf{l}_{11}}-\underbrace{\frac{1} {m}\sum_{j,r}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},y_{i}\mathbf{\mu} \rangle)\beta}_{\mathbf{l}_{12}}, \tag{47}\]

where the first inequality is by Lemma B.2 and the last inequality is by Lemma B.1. Then, we will bound each term in (47) respectively.

For \(\mathbf{l}_{10}\), \(\mathbf{l}_{11}\), \(\mathbf{l}_{12}\), \(\mathbf{l}_{14}\), we have that

\[|\mathbf{l}_{10}| \leq O\big{(}n\sqrt{\log(n/\delta)}/\sqrt{d}\big{)},\;|\mathbf{l }_{11}|\leq O\big{(}n\sqrt{\log(n/\delta)}\cdot\mathrm{SNR}\cdot d^{-1/2}\big{)}, \tag{48}\] \[|\mathbf{l}_{12}| \leq O\big{(}\beta\big{)},\]

For \(j=y_{i}\) and \(r\in S_{i}^{(0)}\), according to Lemma C.3, we have

\[(P-1)\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle \geq(P-1)\langle\mathbf{w}_{j,r}^{(0,0)},\mathbf{\xi}_{i}\rangle+ \overline{\rho}_{j,r,i}^{(t,b)}-5n\sqrt{\frac{\log(4n^{2}/\delta)}{d}}\alpha\] \[\geq 2-\beta-5n\sqrt{\frac{\log(4n^{2}/\delta)}{d}}\alpha\] \[\geq 1.5-\beta>0,\]

where the first inequality is by Lemma C.3; the second inequality is by \(5n\sqrt{\frac{\log(4n^{2}/\delta)}{d}}\leq 0.5\); and the last inequality is by \(\beta<1.5\). Therefore, for \(\mathbf{l}_{9}\), according to the fourth statement of Proposition C.8, we have

\[\mathbf{l}_{9}\geq\frac{1}{m}|\widetilde{S}_{i}^{(t,b)}|20\log(2/\epsilon)\geq 2 \log(2/\epsilon). \tag{49}\]

By plugging (48) and (49) into (47) and according to triangle inequality we have

\[y_{i}\langle\nabla f(\mathbf{W}^{(t,b)},\mathbf{x}_{i}),\mathbf{W}^{*}\rangle \geq\mathbf{l}_{9}-|\mathbf{l}_{10}|-|\mathbf{l}_{11}|-|\mathbf{l}_{12}|-| \mathbf{l}_{14}|\geq\log(2/\epsilon),\]

which completes the proof.

**Lemma C.12**.: _Under Assumption 3.1, for \((0,0)\leq(t,b)\leq(T^{*},0)\), the following result holds._

\[\|\nabla L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})\|_{F}^{2}\leq O(\max\{\|\mathbf{ \mu}\|_{2}^{2},(P-1)^{2}\sigma_{p}^{2}d\})L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)}).\]

Proof.: We first prove that

\[\|\nabla f(\mathbf{W}^{(t,b)},\mathbf{x}_{i})\|_{F}=O(\max\{\|\mathbf{\mu}\|_{2},( P-1)\sigma_{p}\sqrt{d}\}). \tag{50}\]

Without loss of generality, we suppose that \(\widehat{y}_{i}=1\). Then we have that

\[\|\nabla f(\mathbf{W}^{(t,b)},\mathbf{x}_{i})\|_{F} \leq\frac{1}{m}\sum_{j,r}\left\|\big{[}\sigma^{\prime}(\langle \mathbf{w}_{j,r}^{(t,b)},\mathbf{\mu}\rangle)\mathbf{\mu}+(P-1)\sigma^{\prime}(\langle \mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle)\mathbf{\xi}_{i}\big{]}\right\|_{2}\] \[\leq\frac{1}{m}\sum_{j,r}\sigma^{\prime}(\langle\mathbf{w}_{j,r} ^{(t,b)},\mathbf{\mu}\rangle)\|\mathbf{\mu}\|_{2}+\frac{P-1}{m}\sum_{j,r}\sigma^{ \prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi}_{i}\rangle)\|\mathbf{\xi}_{i}\|_ {2}\] \[\leq 4\max\{\|\mathbf{\mu}\|_{2},2(P-1)\sigma_{p}\sqrt{d}\},\]

where the first and second inequalities are by triangle inequality, the third inequality is by Lemma B.1. Then we upper bound the gradient norm \(\|\nabla L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})\|_{F}\) as:

\[\|\nabla L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})\|_{F}^{2}\leq\left[\frac{1} {B}\sum_{i\in\mathcal{I}_{t,b}}\ell^{\prime}\big{(}y_{i}f(\mathbf{W}^{(t,b)}, \mathbf{x}_{i})\big{)}\|\nabla f(\mathbf{W}^{(t,b)},\mathbf{x}_{i})\|_{F} \right]^{2}\]\[\leq\frac{1}{B}\sum_{i\in\mathcal{I}_{t,b}}\frac{1}{m}\sum_{j,r} \Big{(}\Big{|}\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(t,0)},\mathbf{\mu} \rangle\Big{|}+(P-1)\Big{|}\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{( t,0)},\mathbf{\xi}_{i}\rangle\Big{|}\Big{)}\] \[\leq\frac{H\eta(P-1)}{Bm}\|\mu\|_{2}\sigma_{p}\sqrt{2\log(6n/ \delta)}+\frac{H\eta(P-1)^{2}}{Bm}2\sigma_{p}^{2}\sqrt{d\log(6n^{2}/\delta)}\] \[\leq\epsilon,\]

where the first inequality is due to triangle inequality, the second inequality is due to \(|\ell_{i}^{\prime}|\leq 1\), the third inequality is due to triangle inequality and the definition of neural networks, the forth inequality is due to parameter update rule (17) and Lemma B.1, and the fifth inequality is due to Condition 3.1.

**Lemma C.15**.: _Under the same conditions as Theorem 3.2, for all \(T_{1}\leq t\leq T^{*}\), we have \(\max_{j,r,i}|\mathcal{E}_{j,r,i}^{(t,b)}|=\max\big{\{}O\big{(}\sqrt{\log(mn/ \delta)}\cdot\sigma_{0}(P-1)\sigma_{p}\sqrt{d}\big{)},O\big{(}n\sqrt{\log(n/ \delta)}\log(T^{*})/\sqrt{d}\big{)}\big{\}}.\) Besides,_

\[\frac{1}{(s-T_{1})H}\sum_{(T_{1},0)\leq(t,b)<(s,0)}L_{\mathcal{I}_{t,b}}( \mathbf{W}^{(t,b)})\leq\frac{\|\mathbf{W}^{(T_{1},0)}-\mathbf{W}^{*}\|_{F}^{2} }{\eta(s-T_{1})H}+\epsilon\]

_for all \(T_{1}\leq t\leq T^{*}\). Therefore, we can find an iterate with training loss smaller than \(2\epsilon\) within \(T=T_{1}+\left\lfloor\|\mathbf{W}^{(T_{1})}-\mathbf{W}^{*}\|_{F}^{2}/(\eta \epsilon)\right\rfloor=T_{1}+\widetilde{O}(\eta^{-1}\epsilon^{-1}mnd^{-1} \sigma_{p}^{-2})\) iterations._

Proof of Lemma c.15.: Note that \(\max_{j,r,i}|\mathcal{E}_{j,r,i}^{(t)}|=\max\big{\{}O\big{(}\sqrt{\log(mn/ \delta)}\ \cdot\ \sigma_{0}(P\ -1)\sigma_{p}\sqrt{d}\big{)},O\big{(}n\sqrt{\log(n/ \delta)}\log(T^{*})/\sqrt{d}\big{)}\big{\}}\) can be proved in the same way as Lemma C.9.

For any \(T_{1}\leq s\leq T^{*}\), by taking a summation of the inequality in Lemma C.13 and dividing \((s-T_{1})H\) on both sides, we obtain that

\[\frac{1}{(s-T_{1})H}\sum_{(T_{1},0)\leq(t,b)<(s,0)}L_{\mathcal{I}_{t,b}}( \mathbf{W}^{(t,b)})\leq\frac{\|\mathbf{W}^{(T_{1},0)}-\mathbf{W}^{*}\|_{F}^{2} }{\eta(s-T_{1})H}+\epsilon.\]

According to the definition of \(T\), we have

\[\frac{1}{(T-T_{1})H}\sum_{(T_{1},0)\leq(t,b)<(T,0)}L_{\mathcal{I}_{t,b}}( \mathbf{W}^{(t,b)})\leq 2\epsilon.\]

Then there exists an epoch \(T_{1}\leq t\leq T^{*}\) such that

\[\frac{1}{H}\sum_{b=0}^{H-1}L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})\leq 2\epsilon.\]

Thus, according to Lemma C.14, we have

\[L_{S}(\mathbf{W}^{(t,0)})\leq 3\epsilon.\]

**Lemma C.16**.: _Under the same conditions as Theorem 3.2, we have_

\[\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(t,b)}/\gamma_{j^{\prime},r^{\prime}}^{ (t,b)}=\Theta(\mathrm{SNR}^{-2}) \tag{51}\]

_for all \(j,j^{\prime}\in\{\pm 1\}\), \(r,r^{\prime}\in[m]\) and \((T_{2},0)\leq(t,b)\leq(T^{*},0)\)._

Proof.: Now suppose that there exists \((0,0)<(\widetilde{T},0)\leq(T^{*},0)\) such that \(\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(t,0)}/\gamma_{j^{\prime},r^{\prime}} ^{(t,b)}=\Theta(\mathrm{SNR}^{-2})\) for all \((0,0)<(t,0)<(\widetilde{T},0)\). Then for \(\overline{\rho}_{j,r,i}^{(t,b)}\), according to Lemma C.1, we have

\[\gamma_{j,r}^{(t+1,0)}=\gamma_{j,r}^{(t,b)}-\frac{\eta}{Bm}\cdot \sum_{b<H}\bigg{[}\sum_{i\in S_{t}\cap\mathcal{I}_{t,b}}\ell_{i}^{\prime(t,b _{i}^{(t)})}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b_{i}^{(t)})},\widehat {y}_{i}\cdot\mathbf{\mu}\rangle)\] \[\qquad\qquad\qquad\qquad\qquad-\sum_{i\in S_{-}\cap\mathcal{I}_{ t,b}}\ell_{i}^{\prime(t,b_{i}^{(t)})}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b _{i}^{(t)})},\widehat{y}_{i}\cdot\mathbf{\mu}\rangle)\bigg{]}\cdot\|\mathbf{\mu}\|_{2}^{2},\] \[\overline{\rho}_{j,r,i}^{(t+1,0)}=\overline{\rho}_{j,r,i}^{(t,0) }-\frac{\eta(P-1)^{2}}{Bm}\cdot\ell_{i}^{\prime(t,b_{i}^{(t)})}\cdot\sigma^{ \prime}(\langle\mathbf{w}_{j,r}^{(t,b_{i}^{(t)})},\mathbf{\xi}_{i}\rangle)\cdot\| \mathbf{\xi}_{i}\|_{2}^{2}\cdot\mathds{1}(y_{i}=j),\]

It follows that

\[\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(\widetilde{T},0)}\]\[=\sum_{i:y_{i}=j}\overline{\rho}_{j,r,i}^{(\widetilde{T},0)}\] \[=\sum_{i:y_{i}=j}\overline{\rho}_{j,r,i}^{(\widetilde{T}-1,0)}- \frac{\eta(P-1)^{2}}{Bm}\cdot\sum_{i:y_{i}=j}\ell_{i}^{(\widetilde{T}-1,b_{i}^{ (\widetilde{T}-1)})}\cdot\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(\widetilde{ T}-1,b_{i}^{(\widetilde{T}-1)}},\boldsymbol{\xi}_{i}\rangle)\|\boldsymbol{\xi}_{i} \|_{2}^{2}\] \[=\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(\widetilde{T}-1,0)}-\frac {\eta(P-1)^{2}}{Bm}\cdot\sum_{i\in\widetilde{S}_{j,r}^{(\widetilde{T}-1,\xi_{i} ^{(\widetilde{T}-1)})}}\ell_{i}^{(\widetilde{T}-1,b_{i}^{(\widetilde{T}-1)})} \|\boldsymbol{\xi}_{i}\|_{2}^{2}\] \[\geq\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(\widetilde{T}-1)}+ \frac{\eta(P-1)^{2}\sigma_{p}^{2}dH\Phi(-1)}{8m}\cdot\min_{i\in\widetilde{S}_{ j,r}^{(\widetilde{t},\widetilde{b}-1)}\cap\mathcal{I}_{\widetilde{t},\widetilde{b}-1}}| \ell_{i}^{(\widetilde{T}-1,b_{i}^{(\widetilde{T}-1)})}|, \tag{52}\]

where the last equality is by the definition of \(S_{j,r}^{(\widetilde{T}-1)}\) as \(\{i\in[n]:y_{i}=j,\langle\mathbf{w}_{j,r}^{(\widetilde{T}-1)},\boldsymbol{\xi }_{i}\rangle>0\}\); the last inequality is by Lemma B.1 and the fifth statement of Lemma C.8. And

\[\gamma_{j^{\prime},r^{\prime}}^{(\widetilde{T},0)} \leq\gamma_{j^{\prime},r^{\prime}}^{(\widetilde{T}-1,0)}-\frac{ \eta}{Bm}\cdot\sum_{i\in S_{+}}\ell_{i}^{\prime(\widetilde{T}-1,b_{i}^{( \widetilde{T}-1)})}\sigma^{\prime}(\langle\mathbf{w}_{j^{\prime},r^{\prime}}^{ (\widetilde{T}-1,b_{i}^{\widetilde{T}-1}},\widehat{y}_{i}\cdot\boldsymbol{ \mu}\rangle)\cdot\|\boldsymbol{\mu}\|_{2}^{2}\] \[\leq\gamma_{j^{\prime},r^{\prime}}^{(\widetilde{T}-1,0)}+\frac{H \eta\|\boldsymbol{\mu}\|_{2}^{2}}{m}\cdot\max_{i\in S_{+}}|\ell_{i}^{( \widetilde{T}-1)}|. \tag{53}\]

According to the third statement of Lemma C.8, we have \(\max_{i\in S_{+}}|\ell_{i}^{\prime(\widetilde{T}-1,b_{i}^{(\widetilde{T}-1)}) }|\leq C_{2}\min_{i\in S_{j,r}^{(\widetilde{T}-1,b_{i}^{(\widetilde{T}-1)})}} |\ell_{i}^{\prime(\widetilde{T}-1)}|\). Then by combining (52) and (53), we have

\[\frac{\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(\widetilde{T},0)}}{\gamma_{j^{ \prime},r^{\prime}}^{(\widetilde{T},0)}}\geq\min\left\{\frac{\sum_{i=1}^{n} \overline{\rho}_{j,r,i}^{(\widetilde{T}-1,0)}}{\gamma_{j^{\prime},r^{\prime}}^ {(\widetilde{T}-1,0)}},\frac{(P-1)^{2}\sigma_{p}^{2}d}{16C_{2}\|\boldsymbol{ \mu}\|_{2}^{2}}\right\}=\Theta(\mathrm{SNR}^{-2}). \tag{54}\]

On the other hand, we will now show \(\frac{\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(\widetilde{t},0)}}{\gamma_{j^{ \prime},r^{\prime}}^{(\widetilde{t},0)}}\leq\Theta(\mathrm{SNR}^{-2})\) for \(t\geq T_{2}\) by induction. By Lemma B.1 and (52), we have

\[\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(T_{2},0)} \leq\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(T_{2}-1,0)}+\frac{3 \eta(P-1)^{2}\sigma_{p}^{2}dn}{2Bm}\] \[\leq\frac{3\eta(P-1)^{2}\sigma_{p}^{2}dnT_{2}}{2Bm}.\]

And, by (46), we know that at \(t=T_{2}\), we have

\[\gamma_{j^{\prime},r^{\prime}}^{(T_{2},0)}\geq\frac{\eta c_{3}c_{4}CT_{2}n\| \mu\|_{2}^{2}}{8Bm}.\]

Thus,

\[\frac{\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(T_{2},0)}}{\gamma_{j^{\prime},r ^{\prime}}^{(T_{2},0)}}\leq\Theta(\mathrm{SNR}^{-2}).\]

Suppose \(\frac{\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(\widetilde{T},0)}}{\gamma_{j^{ \prime},r^{\prime}}^{(\widetilde{T},0)}}\leq\Theta(\mathrm{SNR}^{-2})\). According to the decomposition, we have:

\[\langle\mathbf{w}_{j,r}^{(T,b)},\widehat{y}_{i}\boldsymbol{\mu}\rangle =\langle\mathbf{w}_{j,r}^{(0,0)},\widehat{y}_{i}\boldsymbol{\mu} \rangle+j\cdot\gamma_{j,r}^{(T,b)}\cdot\widehat{y}_{i}\] \[+\frac{1}{P-1}\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(T,b)}\cdot \|\boldsymbol{\xi}_{i}\|_{2}^{-2}\langle\boldsymbol{\xi}_{i},\widehat{y}_{i} \boldsymbol{\mu}\rangle+\frac{1}{P-1}\sum_{i=1}^{n}\underline{\rho}_{j,r,i}^{(T,b)}\cdot\|\boldsymbol{\xi}_{i}\|_{2}^{-2}\langle\boldsymbol{\xi}_{i},\widehat{y }_{i}\boldsymbol{\mu}\rangle. \tag{55}\]And we have that

\[|\langle\mathbf{w}_{j,r}^{(0,0)},\widehat{y}_{i}\mathbf{\mu}\rangle+ \frac{1}{P-1}\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(T,b)}\cdot\|\mathbf{\xi}_{i}\|_ {2}^{-2}\langle\mathbf{\xi}_{i},\widehat{y}\mathbf{\mu}\rangle+\frac{1}{P-1}\sum_{i=1} ^{n}\rho_{j,r,i}^{(T,b)}\cdot\|\mathbf{\xi}_{i}\|_{2}^{-2}\langle\mathbf{\xi}_{i}, \widehat{y}\mathbf{\mu}\rangle|\] \[\leq\beta/2+|\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(T,b)}|\frac{ 4\|\mu\|_{2}\sqrt{2\log(6n/\delta)}}{\sigma_{p}d(P-1)}\] \[\leq\beta/2+\frac{\Theta(\mathrm{SNR}^{-1})\gamma_{j,r}^{(T,b)}}{ \sqrt{d}}\] \[\leq\gamma_{j,r}^{(T,0)},\]

where the first inequality is due to triangle inequality and Lemma B.1, the second inequality is due to induction hypothesis, and the last inequality is due to Condition 3.1.

Thus, the sign of \(\langle\mathbf{w}_{j,r}^{(T,b)},\widehat{y}_{i}\mathbf{\mu}\rangle\) is persistent through out the epoch. Then, without loss of generality, we suppose \(\langle\mathbf{w}_{j,r}^{(T,b)},\mathbf{\mu}\rangle>0\). Thus, the update rule of \(\gamma\) is:

\[\gamma_{j,r}^{(t,b+1)}\] \[=\gamma_{j,r}^{(t,b)}-\frac{\eta}{Bm}\cdot\bigg{[}\sum_{i\in \mathcal{I}_{T,b}\cap S_{+}\cap S_{1}}\ell_{i}^{\prime(T,b)}-\sum_{i\in \mathcal{I}_{T,b}\cap S_{-}\cap S_{1}}\ell_{i}^{\prime(T,b)}\bigg{]}\|\mathbf{\mu} \|_{2}^{2}\] \[\geq\gamma_{j,r}^{(T,b)}+\frac{\eta}{Bm}\cdot\big{[}\min_{i\in \mathcal{I}_{T,b}}\ell_{i}^{(T,b)}|\mathcal{I}_{T,b}\cap S_{+}\cap S_{1}|-\max _{i\in\mathcal{I}_{T,b}}|\mathcal{I}_{T,b}\cap S_{-}\cap S_{-1}|\big{]}\cdot \|\mathbf{\mu}\|_{2}^{2}.\]

Therefore,

\[\gamma_{j,r}^{(T+1,0)}\geq\gamma_{j,r}^{(T,b)}+\frac{\eta}{Bm}\cdot\big{[}\min \ell_{i}^{(T,b_{i}^{(T)})}|S_{+}\cap S_{1}|-\max\ell_{i}^{(T,b_{i}^{(T)})}|S_ {-}\cap S_{-1}|\big{]}\cdot\|\mathbf{\mu}\|_{2}^{2}. \tag{56}\]

And, by (52), we have

\[\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(T+1,0)}\leq\sum_{i=1}^{n}\overline{ \rho}_{j,r,i}^{(T,0)}+\frac{\eta(P-1)^{2}\sigma_{p}^{2}dH\Phi(-1)}{8m}\cdot \max\Big{|}\ell_{i}^{(T,b_{i}^{(T)})}\Big{|}. \tag{57}\]

Thus, combining (56) and (57), we have

\[\frac{\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(T+1,0)}}{\gamma_{j,r }^{(T+1,0)}}\] \[\leq\max\left\{\frac{\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(T,0 )}}{\gamma_{j,r}^{(T,0)}},\frac{(P-1)^{2}\sigma_{p}^{2}dn\Phi(-1)\cdot\max| \ell_{i}^{\prime(T,b_{i}^{(T)})}|}{8\big{[}\min\ell_{i}^{(T,b_{i}^{(T)})}|S_{ +}\cap S_{1}|-\max\ell_{i}^{(T,b_{i}^{(T)})}|S_{-}\cap S_{-1}|\big{]}\cdot\| \mathbf{\mu}\|_{2}^{2}}\right\}\] \[\leq\Theta(\mathrm{SNR}^{-2}), \tag{58}\]

where the last inequality is due to the induction hypothesis, the third statement of Lemma C.8, and Lemma B.5. Thus, by induction, we have for all \(T_{1}\leq t\leq T^{*}\) that

\[\frac{\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(t,0)}}{\gamma_{j^{\prime},r^{ \prime}}^{(t,0)}}\leq\Theta(\mathrm{SNR}^{-2}).\]

And for \((T_{1},0)\leq(t,b)\leq(T^{*},0)\), we can bound the ratio as follows:

\[\frac{\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(t,b)}}{\gamma_{j^{\prime},r^{ \prime}}^{(t,b)}}\leq\frac{4\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(t,0)}}{ \gamma_{j^{\prime},r^{\prime}}^{(t,0)}}\leq\Theta(\mathrm{SNR}^{-2}),\]

where the first inequality is due to the update rule of \(\overline{\rho}_{j,r,i}^{(t,b)}\) and \(\overline{\rho}_{j,r,i}^{(t,b)}\). Thus, we have completed the proof.

### Test Error Analysis

In this section, we present and prove the exact upper bound and lower bound of test error in Theorem 3.2. Since we have resolved the challenges brought by stochastic mini-batch parameter update, the remaining proof for test error is similar to the counterpart in Kou et al. (2023).

#### c.3.1 Test Error Upper Bound

First, we prove the upper bound of test error in Theorem 3.2 when the training loss converges.

**Theorem C.17** (Second part of Theorem 3.2).: _Under the same conditions as Theorem 3.2, then there exists a large constant \(C_{1}\) such that when \(n\|\mathbf{\mu}\|_{2}^{2}\geq C_{1}(P-1)^{4}\sigma_{p}^{4}d\), for time \(t\) defined in Lemma C.15, we have the test error_

\[\mathbb{P}_{(\mathbf{x},y)\sim\mathcal{D}}\big{(}y\neq\operatorname{sign}(f( \mathbf{W}^{(t,0)},\mathbf{x}))\big{)}\leq p+\exp\bigg{(}-n\|\mathbf{\mu}\|_{2}^{4} /(C_{2}(P-1)^{4}\sigma_{p}^{4}d)\bigg{)},\]

_where \(C_{2}=O(1)\)._

Proof.: The proof is similar to the proof of Theorem E.1 in Kou et al. (2023). The only difference is substituting \(\mathbf{\xi}\) in their proof with \((P-1)\mathbf{\xi}\).

#### c.3.2 Test Error Lower Bound

In this part, we prove the lower bound of the test error in Theorem 3.2. We give two key Lemmas.

**Lemma C.18**.: _For \((T_{1},0)\leq(t,b)<(T^{*},0)\), denote \(g(\mathbf{\xi})=\sum_{j,r}j(P-1)\sigma(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{\xi} \rangle)\). There exists a fixed vector \(\mathbf{v}\) with \(\|\mathbf{v}\|_{2}\leq 0.06\sigma_{p}\) such that_

\[\sum_{j^{\prime}\in\{\pm 1\}}[g(j^{\prime}\mathbf{\xi}+\mathbf{v})-g(j^{\prime}\mathbf{ \xi})]\geq 4C_{6}\max_{j\in\{\pm 1\}}\Big{\{}\sum_{r}\gamma_{j,r}^{(t,b)} \Big{\}}, \tag{59}\]

_for all \(\mathbf{\xi}\in\mathbb{R}^{d}\)._

Proof of Lemma c.18.: The proof is similar to the proof of Lemma 5.8 in Kou et al. (2023). The only difference is substituting \(\mathbf{\xi}\) in their proof with \((P-1)\mathbf{\xi}\).

**Lemma C.19** (Proposition 2.1 in Devroye et al. (2018)).: _The TV distance between \(\mathcal{N}(0,\sigma_{p}^{2}\mathbf{I}_{d})\) and \(\mathcal{N}(\mathbf{v},\sigma_{p}^{2}\mathbf{I}_{d})\) is smaller than \(\|\mathbf{v}\|_{2}/2\sigma_{p}\)._

Then, we can prove the lower bound of the test error.

**Theorem C.20** (Third part of Theorem 3.2).: _Suppose that \(n\|\mathbf{\mu}\|_{2}^{4}\leq C_{3}d(P-1)^{4}\sigma_{p}^{4}\), then we have that \(L_{\mathcal{D}}^{0-1}(\mathbf{W}^{(t,0)})\geq p+0.1\), where \(C_{3}\) is an sufficiently large absolute constant._

Proof.: The proof is similar to the proof of Theorem 4.3 in Kou et al. (2023). The only difference is substituting \(\mathbf{\xi}\) in their proof with \((P-1)\mathbf{\xi}\).

## Appendix D Proofs for SAM

### Noise Memorization Prevention

The following lemma shows the update rule of the neural network

**Lemma D.1**.: _We denote \(\ell_{i}^{\prime(t,b)}=\ell^{\prime}[y_{i}\cdot f(\mathbf{W}^{(t,b)},\mathbf{x }_{i})]\), then the adversarial point of \(\mathbf{W}^{(t,b)}\) is \(\mathbf{W}^{(t,b)}+\mathbf{\widehat{e}}^{(t,b)}\), where_

\[\mathbf{\widehat{e}}^{(t,b)}_{j,r}=\frac{\tau}{m}\frac{\sum_{i\in\mathcal{I}_{t,b }}\sum_{p\in[P]}\ell_{i}^{\prime(t,b)}j\cdot y_{i}\sigma^{\prime}(\langle \mathbf{w}_{j,r}^{(t,b)},\mathbf{x}_{i,p}\rangle)\mathbf{x}_{i,p}}{\|\nabla_{ \mathbf{W}}L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})\|_{F}}.\]

_Then the training update rule of the parameter is_

\[\mathbf{w}_{j,r}^{(t+1,b)}=\mathbf{w}_{j,r}^{(t,b)}-\frac{\eta}{Bm}\sum_{i\in \mathcal{I}_{t,b}}\sum_{p\in[P]}\ell_{i}^{\prime(t,b)}\sigma^{\prime}(\langle \mathbf{w}_{j,r}^{(t,b)}+\mathbf{\widehat{e}}_{j,r}^{(t,b)},\mathbf{x}_{i,p}\rangle )j\cdot\mathbf{x}_{i,p}\]\[=\mathbf{w}_{j,r}^{(t,b)}-\frac{\eta}{Bm}\sum_{i\in\mathcal{I}_{t,b}} \sum_{p\in[P]}\ell_{i}^{\prime(t,b)}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{x}_{i,p}\rangle+\langle\widehat{\mathbf{\epsilon}}_{t,j,r},\mathbf{x}_ {i,p}\rangle)j\cdot\mathbf{x}_{i,p}\] \[=\mathbf{w}_{j,r}^{(t,b)}-\frac{\eta}{Bm}\sum_{i\in\mathcal{I}_{t,b}}\ell_{i}^{\prime(t,b)}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},y \mathbf{\mu}\rangle+\langle\widehat{\mathbf{\epsilon}}_{t,j,r},y\mathbf{\mu}\rangle)j\mathbf{\mu}\] \[\qquad\underbrace{-\frac{\eta(P-1)}{Bm}\sum_{i\in\mathcal{I}_{t,b }}\ell_{i}^{\prime(t,b)}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\mathbf{ \xi}_{i}\rangle+\langle\widehat{\mathbf{\epsilon}}_{j,r}^{(t,b)},\mathbf{\xi}_{i} \rangle)jy_{i}\mathbf{\xi}_{i}}_{\text{NoiseTerm}}.\]

We will show that the noise term will be small if we train with the SAM algorithm. We consider the first stage where \(t\leq T_{1}\) where \(T_{1}=mB/(12n\eta||\mathbf{\mu}||_{2}^{2})\). Then, the following property holds.

**Proposition D.2**.: _Under Assumption 3.1, for \(0\leq t\leq T_{1}\), we have that_

\[\gamma_{j,r}^{(0,0)},\overline{\rho}_{j,r,i}^{(0,0)},\rho_{j,r,i} ^{(0,0)}=0, \tag{60}\] \[0\leq\gamma_{j,r}^{(t,b)}\leq 1/12,\] (61) \[0\leq\overline{\rho}_{j,r,i}^{(t,b)}\leq 1/12,\] (62) \[0\geq\underline{\rho}_{j,r,i}^{(t,b)}\geq-\beta-10\sqrt{\frac{ \log(6n^{2}/\delta)}{d}}n. \tag{63}\]

_Besides, \(\gamma_{j,r}^{(T_{1},0)}=\Omega(1)\)._

**Lemma D.3**.: _Under Assumption 3.1, suppose (27), (28) and (29) hold at iteration \(t\). Then, for all \(r\in[m]\), \(j\in\{\pm 1\}\) and \(i\in[n]\),_

\[\left|\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(0,0)}, \mathbf{\mu}\rangle-j\cdot\gamma_{j,r}^{(t,b)}\right|\leq\mathrm{SNR}\sqrt{\frac{3 2\log(6n/\delta)}{d}}n\alpha, \tag{64}\] \[\left|\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(0,0)}, \mathbf{\xi}_{i}\rangle-\frac{1}{P-1}\underline{\rho}_{j,r,i}^{(t,b)}\right|\leq \frac{5}{P-1}\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha,\,j\neq y_{i},\] (65) \[\left|\langle\mathbf{w}_{j,r}^{(t,b)}-\mathbf{w}_{j,r}^{(0,0)}, \mathbf{\xi}_{i}\rangle-\frac{1}{P-1}\overline{\rho}_{j,r,i}^{(t,b)}\right|\leq \frac{5}{P-1}\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha,\,j=y_{i}. \tag{66}\]

Proof of Lemma d.3.: Notice that \(1/12<\alpha\), if the condition (61), (62), (63) holds, (27), (28) and (29) also holds. Therefore, by Lemma C.3, we know that Lemma D.3 also holds.

**Lemma D.4**.: _Under Assumption 3.1, suppose (61), (62), (63) hold at iteration \(t,b\). Then, for all \(j\in\{\pm 1\}\) and \(i\in[n]\), \(F_{j}(\mathbf{W}_{j}^{(t,b)},\mathbf{x}_{i})\leq 0.5\). Therefore \(-0.3\geq\ell_{i}^{\prime}\geq-0.7\)._

Proof.: Notice that \(1/12<\alpha\), if the condition (61), (62), (63) holds, (27), (28) and (29) also holds. Therefore, by Lemma C.4, we know that for all \(j\neq y_{i}\) and \(i\in[n]\), \(F_{j}(\mathbf{W}_{j}^{(t,b)},\mathbf{x}_{i})\leq 0.5\). Next we will show that for \(j=y_{i}\), \(F_{j}(\mathbf{W}_{j}^{(t,b)},\mathbf{x}_{i})\leq 0.5\) also holds.

According to Lemma D.3, we have

\[F_{j}(\mathbf{W}_{j}^{(t,b)},\mathbf{x}_{i}) =\frac{1}{m}\sum_{r=1}^{m}[\sigma(\langle\mathbf{w}_{j,r}^{(t,b)},y_{i}\mathbf{\mu}\rangle)+(P-1)\sigma(\langle\mathbf{w}_{j,r}^{(t)},\mathbf{\xi}_{i }\rangle)]\] \[\leq 2\max\{|\langle\mathbf{w}_{j,r}^{(t,b)},y_{i}\mathbf{\mu}\rangle |,(P-1)|\langle\mathbf{w}_{j,r}^{(t)},\mathbf{\xi}_{i}\rangle|\}\] \[\leq 6\max\left\{|\langle\mathbf{w}_{j,r}^{(0)},\widehat{y}_{i}\bm {\mu}\rangle|,(P-1)|\langle\mathbf{w}_{j,r}^{(0)},\mathbf{\xi}_{i}\rangle|, \mathrm{SNR}\sqrt{\frac{32\log(6n/\delta)}{d}}n\alpha,\right.\] \[\qquad\left.5\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha,|\gamma_{ j,r}^{(t,b)}|,|\underline{\rho}_{j,r,i}^{(t,b)}\right\}\] \[\leq 6\max\left\{\beta,\mathrm{SNR}\sqrt{\frac{32\log(6n/\delta)}{ d}}n\alpha,5\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha,|\gamma_{j,r}^{(t,b)}|,| \overline{\rho}_{j,r,i}^{(t,b)}|\right\}\]\[\leq 0.5,\]

where the second inequality is by (64), (65) and (66); the third inequality is due to the definition of \(\beta\); the last inequality is by (25), (61), (62).

Since \(F_{j}(\mathbf{W}_{j}^{(t,b)},\mathbf{x}_{i})\in[0,0.5]\) we know that

\[-0.3\geq-\frac{1}{1+\exp(0.5)}\geq\ell_{i}^{\prime}\geq-\frac{1}{1+\exp(-0.5)} \geq-0.7.\]

Based on the previous foundation lemmas, we can provide the key lemma of SAM which is different from the dynamic of SGD.

**Lemma D.5**.: _Under Assumption 3.1, suppose (61), (62) and (63) hold at iteration \(t,b\). We have that if \(\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\xi}_{k}\rangle\geq 0\), \(k\in\mathcal{I}_{t,b}\) and \(j=y_{k}\), then \(\langle\mathbf{w}_{j,r}^{(t,b)}+\boldsymbol{\hat{\epsilon}}_{j,r}^{(t,b)}, \boldsymbol{\xi}_{k}\rangle<0\)._

Proof.: We first prove that there for \(t\leq T_{1}\), there exists a constant \(C_{2}\) such that

\[\|\nabla_{\mathbf{W}}L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})\|_{F}\leq C_{2 }P\sigma_{p}\sqrt{d/B}.\]

Recall that

\[L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})=\frac{1}{B}\sum_{i\in\mathcal{I}_{t,b}}\ell(y_{i}f(\mathbf{W}^{(t,b)},x_{i})),\]

we have

\[\nabla_{\mathbf{w}_{j,r}}L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)}) =\frac{1}{B}\sum_{i\in\mathcal{I}_{t,b}}\nabla_{\mathbf{w}_{j,r}} \ell(y_{i}f(\mathbf{W}^{(t,b)},\mathbf{x}_{i}))\] \[=\frac{1}{B}\sum_{i\in\mathcal{I}_{t,b}}y_{i}\ell^{\prime}(y_{i}f( \mathbf{W}^{(t,b)},\mathbf{x}_{i}))\nabla_{\mathbf{w}_{j,r}}f(\mathbf{W}^{(t, b)},\mathbf{x}_{i})\] \[=\frac{1}{Bm}\sum_{i\in\mathcal{I}_{t,b}}y_{i}\ell_{i}^{\prime(t, b)}[\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\mu}\rangle) \cdot\boldsymbol{\mu}+\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)}, \boldsymbol{\xi}_{i}\rangle)\cdot(P-1)\boldsymbol{\xi}_{i}].\]

We have

\[\|\nabla_{\mathbf{w}_{j,r}}L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b )})\|_{2}\] \[\leq\frac{1}{Bm}\bigg{\|}\sum_{i\in\mathcal{I}_{t,b}}|\ell_{i}^{ \prime(t,b)}|\cdot\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t)},\boldsymbol{ \mu}\rangle)\cdot\boldsymbol{\mu}\bigg{\|}_{2}+\frac{1}{Bm}\bigg{\|}\sum_{i \in\mathcal{I}_{t,b}}|\ell_{i}^{\prime(t,b)}|\cdot\sigma^{\prime}(\langle \mathbf{w}_{j,r}^{(t)},\boldsymbol{\xi}_{i}\rangle)\cdot(P-1)\boldsymbol{\xi }_{i}\bigg{\|}_{2}\] \[\leq 0.7m^{-1}\|\boldsymbol{\mu}\|_{2}+1.4(P-1)m^{-1}\sigma_{p} \sqrt{d/B}\] \[\leq 2Pm^{-1}\sigma_{p}\sqrt{d/B},\]

and

\[\|\nabla_{\mathbf{W}}L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})\|_{F}^{2}=\sum _{j,r}\|\nabla_{\mathbf{w}_{j,r}}L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})\|_{ 2}^{2}\leq 2m(2Pm^{-1}\sigma_{p}\sqrt{d/B})^{2},\]

leading to

\[\|\nabla_{\mathbf{W}}L_{\mathcal{I}_{t,b}}(\mathbf{W}^{(t,b)})\|_{F}\leq 2 \sqrt{2}P\sigma_{p}\sqrt{d/Bm}.\]

From Lemma D.1, we have

\[\langle\boldsymbol{\hat{\epsilon}}_{j,r}^{(t,b)},\boldsymbol{\xi} _{k}\rangle =\frac{\tau}{mB}\|\nabla_{\mathbf{W}}L_{\mathcal{I}_{t,b}}( \mathbf{W}^{(t,b)})\|_{F}^{-1}\sum_{i\in\mathcal{I}_{t,b}}\sum_{p\in[P]}\ell _{i}^{\prime(t)}j\cdot y_{i}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t)}, \mathbf{x}_{i,p}\rangle)\langle\mathbf{x}_{i,p},\boldsymbol{\xi}_{k}\rangle\] \[=\frac{\tau}{mB}\|\nabla_{\mathbf{W}}L_{\mathcal{I}_{t,b}}( \mathbf{W}^{(t,b)})\|_{F}^{-1}\cdot\bigg{(}\sum_{i\in\mathcal{I}_{t,b},i\neq k }\ell_{i}^{\prime(t,b)}jy_{i}\cdot(P-1)\sigma^{\prime}(\langle\mathbf{w}_{j, r}^{(t,b)},\boldsymbol{\xi}_{i}\rangle)\langle\boldsymbol{\xi}_{i},\boldsymbol{\xi}_{k}\rangle\] \[\qquad+\ell_{k}^{\prime(t)}jy_{k}\cdot(P-1)\sigma^{\prime}(\langle \mathbf{w}_{j,r}^{(t)},\boldsymbol{\xi}_{k}\rangle)\langle\boldsymbol{\xi}_{k},\boldsymbol{\xi}_{k}\rangle\]\[\gamma_{j,r}^{(T,\widetilde{b})}\leq T_{1}\cdot(n/B)\cdot\frac{\eta}{m}\| \boldsymbol{\mu}\|_{2}^{2}\leq 1/12.\]Second, by Lemma D.6, we know that (62) holds for \((\widetilde{T}-1,\widetilde{b})\).

Last, we need to prove that (63) holds \((\widetilde{T}-1,\widetilde{b})\). The proof is similar to previous proof without SAM.

When \(\rho_{j,r,k}^{(\widetilde{T}-1,\widetilde{b}-1)}<-0.5(P-1)\beta-6\sqrt{\frac{ \log(6n^{2}/\delta)}{d}}n\alpha\), by (31), we have

\[\langle\mathbf{w}_{j,r}^{(\widetilde{T}-1,\widetilde{b}-1)}, \boldsymbol{\xi}_{k}\rangle <\langle\mathbf{w}_{j,r}^{(0,0)},\boldsymbol{\xi}_{k}\rangle+ \frac{1}{P-1}\rho_{j,r,k}^{(\widetilde{T}-1,\widetilde{b}-1)}+\frac{5}{P-1} \sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha\] \[\leq-\frac{1}{P-1}\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha,\]

and we have

\[\langle\boldsymbol{\widehat{\epsilon}}_{j,r}^{(\widetilde{T}-1, \widetilde{b}-1)},\boldsymbol{\xi}_{i}\rangle =\frac{\tau}{mB}\|\nabla_{\mathbf{W}}L_{\mathcal{I}_{\widetilde{T }-1,\widetilde{b}-1}}(\mathbf{W}^{(\widetilde{T}-1,\widetilde{b}-1)})\|_{F}^{ -1}\sum_{i\in\mathcal{I}_{\widetilde{T}-1,\widetilde{b}-1}}\sum_{p\in[P]}\ell _{i}^{\prime(\widetilde{T}-1,\widetilde{b}-1)}j\cdot y_{i}\] \[\qquad\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(\widetilde{T}-1, \widetilde{b}-1)},\boldsymbol{x}_{i,p}\rangle)\langle\mathbf{x}_{i,p}, \boldsymbol{\xi}_{k}\rangle\] \[=\frac{\tau}{mB}\|\nabla_{\mathbf{W}}L_{\mathcal{I}_{\widetilde{T }-1,\widetilde{b}-1}}(\mathbf{W}^{(\widetilde{T}-1,\widetilde{b}-1)})\|_{F}^{ -1}\cdot\bigg{(}\sum_{i\in\mathcal{I}_{\widetilde{T}-1,\widetilde{b}-1},i\neq k }\ell_{i}^{\prime(\widetilde{T}-1,\widetilde{b}-1)}j\cdot y_{i}\] \[\qquad(P-1)\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(\widetilde{ T}-1,\widetilde{b}-1)},\boldsymbol{\xi}_{i}\rangle)\langle\boldsymbol{\xi}_{i}, \boldsymbol{\xi}_{k}\rangle+\ell_{k}^{\prime(t)}jy_{k}\cdot(P-1)\sigma^{\prime }(\langle\mathbf{w}_{j,r}^{(t)},\boldsymbol{\xi}_{k}\rangle)\langle\boldsymbol {\xi}_{k},\boldsymbol{\xi}_{k}\rangle\] \[\qquad+\sum_{i\in\mathcal{I}_{\widetilde{T}-1,\widetilde{b}-1}} \ell_{i}^{\prime(\widetilde{T}-1,\widetilde{b}-1)}j\cdot\sigma^{\prime}( \langle\mathbf{w}_{j,r}^{(\widetilde{T}-1,\widetilde{b}-1)},y_{i}\boldsymbol{ \mu}\rangle\langle\boldsymbol{\mu},\boldsymbol{\xi}_{k}\rangle\bigg{)}\] \[\leq\frac{\tau}{mC_{2}P\sigma_{p}\sqrt{Bd}}\bigg{[}0.8B(P-1) \sigma_{P}^{2}\sqrt{d\log(6n^{2}/\delta)}+0.4B\sigma_{P}\|\boldsymbol{\mu}\|_{ 2}\sqrt{2\log(6n^{2}/\delta)}\bigg{]}\] \[\leq C_{4}\frac{\tau\sqrt{B}\sigma_{p}\sqrt{\log(6n^{2}/\delta)} }{m}\] \[=C_{4}\frac{B\sqrt{\log(6n^{2}/\delta)}}{C_{3}P\sqrt{d}}\] \[\leq\frac{1}{P}\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha,\]

and thus \(\langle\mathbf{w}_{j,r}^{(\widetilde{T}-1,\widetilde{b}-1)}+\boldsymbol{\widehat {\epsilon}}_{j,r}^{(\widetilde{T}-1,\widetilde{b}-1)},\boldsymbol{\xi}_{i} \rangle<0\) which leads to

\[\rho_{j,r,i}^{(\widetilde{T}-1,\widetilde{b})} =\rho_{j,r,i}^{(\widetilde{T}-1,\widetilde{b}-1)}+\frac{\eta(P- 1)^{2}}{Bm}\cdot\ell_{i}^{\prime(\widetilde{T}-1,\widetilde{b}-1)}\cdot\sigma^ {\prime}(\langle\mathbf{w}_{j,r}^{(\widetilde{T}-1,\widetilde{b}-1)}, \boldsymbol{\xi}_{i}\rangle)\cdot\|\boldsymbol{\xi}_{i}\|_{2}^{2}.\] \[\quad\mathds{1}(y_{i}=-j)\,\mathds{1}(i\in\mathcal{I}_{\widetilde {T}-1,\widetilde{b}-1})\] \[=\rho_{j,r,i}^{(\widetilde{T}-1,\widetilde{b}-1)}.\]

Therefore, we have

\[\rho_{j,r,i}^{(\widetilde{T}-1,\widetilde{b})}=\rho_{j,r,i}^{(\widetilde{T}-1, \widetilde{b}-1)}\geq-(P-1)\beta-5P\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha.\]

When \(\rho_{j,r,i}^{(\widetilde{T}-1,\widetilde{b}-1)}\geq-0.5(P-1)\beta-5\sqrt{\frac{ \log(6n^{2}/\delta)}{d}}n\alpha\), we have that

\[\rho_{j,r,i}^{(\widetilde{T}-1,\widetilde{b})} \geq\rho_{j,r,i}^{(\widetilde{T}-1,\widetilde{b}-1)}+\frac{\eta(P- 1)^{2}}{Bm}\cdot\ell_{i}^{\prime(\widetilde{T}-1,\widetilde{b}-1)}\cdot\| \boldsymbol{\xi}_{i}\|_{2}^{2}\] \[\geq\rho_{j,r,i}^{(\widetilde{T}-1,\widetilde{b}-1)}-\frac{0.4\eta (P-1)^{2}}{Bm}\cdot 2d\sigma_{p}^{2}\] \[\geq-(P-1)\beta-5P\sqrt{\frac{\log(6n^{2}/\delta)}{d}}n\alpha.\]Therefore, the induction is completed, and thus Proposition D.2 holds.

Next, we will prove that \(\gamma_{j,r}^{(t)}\) can achieve \(\Omega(1)\) after \(T_{1}=mB/(12n\eta\|\mathbf{\mu}\|_{2}^{2})\) iterations. By Lemma B.6, we know that there exists \(c_{3}\cdot T_{1}\) epochs such that at least \(c_{4}\cdot H\) batches in these epochs, satisfy

\[|S_{+}\cap S_{y}\cap\mathcal{I}_{t,b}|\in\left[\frac{B}{4},\frac{3B}{4}\right]\]

for both \(y=+1\) and \(y=-1\). For SAM, we have the following update rule for \(\gamma_{j,r}^{(t,b)}\):

\[\gamma_{j,r}^{(t,b+1)} =\gamma_{j,r}^{(t,b)}-\frac{\eta}{Bm}\sum_{i\in\mathcal{I}_{t,b} \cap S_{+}}\ell_{i}^{\prime(t,b)}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t, b)}+\widehat{\mathbf{\epsilon}}_{j,r}^{(t,b)},y_{i}\cdot\mathbf{\mu}\rangle)\cdot\|\mathbf{ \mu}\|_{2}^{2}\] \[\qquad+\frac{\eta}{Bm}\sum_{i\in\mathcal{I}_{t,b}\cap S_{-}}\ell _{i}^{\prime(t,b)}\sigma^{\prime}(\langle\mathbf{w}_{j,r}^{(t,b)}+\widehat{\bm {\epsilon}}_{j,r}^{(t,b)},y_{i}\cdot\mathbf{\mu}\rangle)\cdot\|\mathbf{\mu}\|_{2}^{2}.\]

If \(\langle\mathbf{w}_{j,r}^{(t,b)}+\widehat{\mathbf{\epsilon}}_{j,r}^{(t,b)},\mathbf{\mu} \rangle\geq 0\), we have

\[\gamma_{j,r}^{(t,b+1)} =\gamma_{j,r}^{(t,b)}-\frac{\eta}{Bm}\cdot\bigg{[}\sum_{i\in \mathcal{I}_{t,b}\cap S_{+}\cap S_{1}}\ell_{i}^{\prime(t)}-\sum_{i\in\mathcal{ I}_{t,b}\cap S_{+}\cap S_{-1}}\ell_{i}^{\prime(t)}\bigg{]}\|\mathbf{\mu}\|_{2}^{2}\] \[\geq\gamma_{j,r}^{(t,b)}+\frac{\eta}{Bm}\cdot\big{(}0.3|\mathcal{ I}_{t,b}\cap S_{+}\cap S_{1}|-0.7|\mathcal{I}_{t,b}\cap S_{+}\cap S_{-1}| \big{)}\cdot\|\mathbf{\mu}\|_{2}^{2}.\]

If \(\langle\mathbf{w}_{j,r}^{(t,b)}+\widehat{\mathbf{\epsilon}}_{j,r}^{(t,b)},\mathbf{\mu} \rangle<0\), we have

\[\gamma_{j,r}^{(t,b+1)} =\gamma_{j,r}^{(t,b)}+\frac{\eta}{Bm}\cdot\bigg{[}\sum_{i\in \mathcal{I}_{t,b}\cap S_{+}\cap S_{-1}}\ell_{i}^{(t)}-\sum_{i\in\mathcal{I}_{t,b}\cap S_{+}\cap S_{1}}\ell_{i}^{\prime(t)}\bigg{]}\|\mathbf{\mu}\|_{2}^{2}\] \[\geq\gamma_{j,r}^{(t,b)}+\frac{\eta}{Bm}\cdot\big{(}0.3|\mathcal{ I}_{t,b}\cap S_{+}\cap S_{-1}|-0.7|\mathcal{I}_{t,b}\cap S_{+}\cap S_{1}| \big{)}\cdot\|\mathbf{\mu}\|_{2}^{2}.\]

Therefore, we have

\[\gamma_{j,r}^{(T_{1},0)} \geq\frac{\eta}{Bm}(0.3\cdot c_{3}T_{1}\cdot c_{4}H\cdot 0.25B-0.7T_ {1}nq)\|\mathbf{\mu}\|_{2}^{2}\] \[=\frac{\eta}{Bm}(0.075c_{3}c_{4}T_{1}n-0.7T_{1}nq)\|\mathbf{\mu}\|_{2} ^{2}\] \[\geq\frac{\eta}{16Bm}c_{3}c_{4}T_{1}n\|\mathbf{\mu}\|_{2}^{2}\] \[=\frac{c_{3}c_{4}}{192}=\Omega(1).\]

**Lemma D.7**.: _Suppose Condition 3.1 holds. Then we have that \(\big{\|}\mathbf{w}_{j,r}^{(T_{1},0)}\big{\|}_{2}=\Theta(\sigma_{0}\sqrt{d})\) and_

\[\langle\mathbf{w}_{j,r}^{(T_{1},0)},j\mathbf{\mu}\rangle=\Omega(1),\] \[\langle\mathbf{w}_{-j,r}^{(T_{1},0)},j\mathbf{\mu}\rangle=-\Omega(1),\] \[\widehat{\beta}:=2\max_{i,j,r}\{|\langle\mathbf{w}_{j,r}^{(T_{1 },0)},\mathbf{\mu}\rangle|,(P-1)|\langle\mathbf{w}_{j,r}^{(T_{1},0)},\mathbf{\xi}_{i} \rangle|\}=O(1).\]

_Besides, for \(S_{i}^{(t,b)}\) and \(S_{j,r}^{(t,b)}\) defined in Lemma B.3 and B.4, we have that_

\[|S_{i}^{(T_{1},0)}|=\Omega(m),\,\forall i\in[n]\] \[|S_{j,r}^{(T_{1})}|=\Omega(n),\,\forall j\in\{\pm 1\},r\in[m].\]

### Test Error Analysis

Proof of Theorem 4.1.: Recall that

\[\mathbf{w}_{j,r}^{(t,b)}=\mathbf{w}_{j,r}^{(0,0)}+j\cdot\gamma_{j,r}^{(t,b)}\cdot \|\boldsymbol{\mu}\|_{2}^{-2}\cdot\boldsymbol{\mu}+\frac{1}{P-1}\sum_{i=1}^{n} \rho_{j,r,i}^{(t,b)}\cdot\|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot\boldsymbol{\xi }_{i},\]

by triangle inequality we have

\[\left|\left\|\mathbf{w}_{j,r}^{(T_{1},0)}\right\|_{2}-\left\| \mathbf{w}_{j,r}^{(0,0)}\right\|_{2}\right| \leq|\gamma_{j,r}^{(t,b)}|\cdot\|\boldsymbol{\mu}\|_{2}^{-1}+ \frac{1}{P-1}\bigg{\|}\sum_{i=1}^{n}|\rho_{j,r,i}^{(t,b)}|\cdot\|\boldsymbol{ \xi}_{i}\|_{2}^{-2}\cdot\boldsymbol{\xi}_{i}\bigg{\|}_{2}\] \[\leq\frac{1}{12}\|\boldsymbol{\mu}\|_{2}^{-1}+\frac{\sqrt{n}}{12 (P-1)}(\sigma_{p}^{2}d/2)^{-1/2}\] \[\leq\frac{1}{6}\|\boldsymbol{\mu}\|_{2}^{-1}.\]

By the condition on \(\sigma_{0}\) and Lemma B.2, we have

\[\left\|\mathbf{w}_{j,r}^{(T_{1},0)}\right\|_{2}=\Theta\big{(}\big{\|}\mathbf{w }_{j,r}^{(0,0)}\big{\|}_{2}\big{)}=\Theta(\sigma_{0}\sqrt{d}).\]

By taking the inner product with \(\boldsymbol{\mu}\) and \(\boldsymbol{\xi}_{i}\), we can get

\[\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\mu}\rangle=\langle\mathbf{w}_{j,r }^{(0,0)},\boldsymbol{\mu}\rangle+j\cdot\gamma_{j,r}^{(t,b)}+\frac{1}{P-1}\sum_ {i=1}^{n}\rho_{j,r,i}^{(t,b)}\cdot\|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot \langle\boldsymbol{\xi}_{i},\boldsymbol{\mu}\rangle,\]

and

\[\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\xi}_{i}\rangle =\langle\mathbf{w}_{j,r}^{(0,0)},\boldsymbol{\xi}_{i}\rangle+j \cdot\gamma_{j,r}^{(t,b)}\cdot\|\boldsymbol{\mu}\|_{2}^{-2}\cdot\langle \boldsymbol{\mu},\boldsymbol{\xi}_{i}\rangle+\frac{1}{P-1}\sum_{i^{\prime}=1} ^{n}\rho_{j,r,i^{\prime}}^{(t,b)}\cdot\|\boldsymbol{\xi}_{i^{\prime}}\|_{2}^ {-2}\cdot\langle\boldsymbol{\xi}_{i^{\prime}},\boldsymbol{\xi}_{i}\rangle\] \[=\langle\mathbf{w}_{j,r}^{(0,0)},\boldsymbol{\xi}_{i}\rangle+j \cdot\gamma_{j,r}^{(t,b)}\cdot\|\boldsymbol{\mu}\|_{2}^{-2}\cdot\langle \boldsymbol{\mu},\boldsymbol{\xi}_{i}\rangle+\frac{1}{P-1}\rho_{j,r,i}^{(t,b)}\] \[\qquad+\frac{1}{P-1}\sum_{i\neq i^{\prime}}\rho_{j,r,i^{\prime}} ^{(t,b)}\cdot\|\boldsymbol{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot\langle \boldsymbol{\xi}_{i^{\prime}},\boldsymbol{\xi}_{i}\rangle.\]

Then, we have

\[\langle\mathbf{w}_{j,r}^{(T_{1},0)},j\boldsymbol{\mu}\rangle =\langle\mathbf{w}_{j,r}^{(0,0)},j\boldsymbol{\mu}\rangle+\gamma _{j,r}^{(T_{1},0)}+\frac{1}{P-1}\sum_{i=1}^{n}\rho_{j,r,i}^{(T_{1},0)}\cdot \|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot\langle\boldsymbol{\xi}_{i},j\boldsymbol {\mu}\rangle\] \[\geq\gamma_{j,r}^{(T_{1},0)}-|\langle\mathbf{w}_{j,r}^{(0,0)}, \boldsymbol{\mu}\rangle|-\frac{1}{P-1}\sum_{i=1}^{n}|\rho_{j,r,i}^{(T_{1},0)}| \cdot\|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot|\langle\boldsymbol{\xi}_{i}, \boldsymbol{\mu}\rangle|\] \[\geq\gamma_{j,r}^{(T_{1},0)}-\sqrt{2\log(12m/\delta)}\cdot\sigma_ {0}\|\boldsymbol{\mu}\|_{2}-\frac{n}{12(P-1)}(\sigma_{0}^{2}d/2)^{-1}\| \boldsymbol{\mu}\|_{2}\sigma_{p}\cdot\sqrt{2\log(6n/\delta)}\] \[\geq\frac{1}{2}\gamma_{j,r}^{(T_{1},0)},\]

and

\[\langle\mathbf{w}_{-j,r}^{(T_{1},0)},j\boldsymbol{\mu}\rangle =\langle\mathbf{w}_{-j,r}^{(0,0)},j\boldsymbol{\mu}\rangle-\gamma _{-j,r}^{(T_{1},0)}-\frac{1}{P-1}\sum_{i=1}^{n}\rho_{-j,r,i}^{(T_{1},0)}\cdot \|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot\langle\boldsymbol{\xi}_{i},j\boldsymbol {\mu}\rangle\] \[\leq-\gamma_{-j,r}^{(T_{1},0)}+|\langle\mathbf{w}_{-j,r}^{(0,0)}, \boldsymbol{\mu}\rangle|+\frac{1}{P-1}\sum_{i=1}^{n}|\rho_{-j,r,i}^{(T_{1},0)}| \cdot\|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot|\langle\boldsymbol{\xi}_{i}, \boldsymbol{\mu}\rangle|\] \[\leq-\gamma_{-j,r}^{(T_{1},0)}+\sqrt{2\log(12m/\delta)}\cdot \sigma_{0}\|\boldsymbol{\mu}\|_{2}+\frac{n}{12(P-1)}(\sigma_{0}^{2}d/2)^{-1}\| \boldsymbol{\mu}\|_{2}\sigma_{p}\cdot\sqrt{2\log(6n/\delta)}\] \[\leq-\frac{1}{2}\gamma_{j,r}^{(T_{1},0)},\]where the last inequality is by the condition on \(\sigma_{0}\) and \(\gamma_{j,r}^{(T_{1},0)}=\Omega(1)\). Thus, it follows that

\[\langle\mathbf{w}_{j,r}^{(T_{1},0)},j\boldsymbol{\mu}\rangle=\Omega(1),\;\langle \mathbf{w}_{-j,r}^{(T_{1},0)},j\boldsymbol{\mu}\rangle=-\Omega(1).\]

By triangle inequality, we have

\[|\langle\mathbf{w}_{j,r}^{(T_{1},0)},\boldsymbol{\mu}\rangle| \leq|\langle\mathbf{w}_{j,r}^{(0,0)},\boldsymbol{\mu}\rangle|+| \gamma_{j,r}^{(T_{1},0)}|+\frac{1}{P-1}\sum_{i=1}^{n}|\rho_{j,r,i}^{(t,b)}|\cdot \|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot|\langle\boldsymbol{\xi}_{i},\boldsymbol{ \mu}\rangle|\] \[\leq\frac{1}{2}\beta+\frac{1}{12}+\frac{n}{P-1}\cdot\frac{1}{12} (\sigma_{p}^{2}d/2)^{-1}\cdot\|\boldsymbol{\mu}\|_{2}\sigma_{p}\cdot\sqrt{2 \log(6n/\delta)}\] \[=\frac{1}{2}\beta+\frac{1}{12}+\frac{n}{6(P-1)}\|\boldsymbol{\mu }\|_{2}\sqrt{2\log(6n/\delta)}/(\sigma_{p}d)\] \[\leq\frac{1}{6},\]

and

\[|\langle\mathbf{w}_{j,r}^{(T_{1},0)},\boldsymbol{\xi}_{i}\rangle| \leq|\langle\mathbf{w}_{j,r}^{(0,0)},\boldsymbol{\xi}_{i}\rangle| +|\gamma_{j,r}^{(T_{1},0)}|\cdot\|\boldsymbol{\mu}\|_{2}^{-2}\cdot|\langle \boldsymbol{\mu},\boldsymbol{\xi}_{i}\rangle|+\frac{1}{P-1}|\rho_{j,r,i}^{(T_ {1},0)}|\] \[\qquad+\frac{1}{P-1}\sum_{i\neq i^{\prime}}|\rho_{j,r,i^{\prime}}^ {(T_{1},0)}|\cdot\|\boldsymbol{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot|\langle \boldsymbol{\xi}_{i^{\prime}},\boldsymbol{\xi}_{i}\rangle|\] \[\leq\frac{1}{2}\beta+\frac{1}{12}\|\boldsymbol{\mu}\|_{2}^{-1} \sigma_{p}\cdot\sqrt{2\log(6n/\delta)}+\frac{1}{12(P-1)}\] \[\qquad+\frac{n}{12(P-1)}(\sigma_{p}^{2}d/2)^{-1}2\sigma_{p}^{2} \cdot\sqrt{d\log(6n^{2}/\delta)}\] \[\leq\frac{1}{2}\beta+\frac{1}{12(P-1)}+\frac{1}{6}\|\boldsymbol{ \mu}\|_{2}^{-1}\sigma_{p}\cdot\sqrt{\log(6n/\delta)}\] \[\leq\frac{1}{6}.\]

This leads to

\[\widehat{\beta}:=2\max_{i,j,r}\{|\langle\mathbf{w}_{j,r}^{(T_{1},0)}, \boldsymbol{\mu}\rangle|,(P-1)|\langle\mathbf{w}_{j,r}^{(T_{1},0)},\boldsymbol {\xi}_{i}\rangle|\}=O(1).\]

In addition, we also have for \(t\leq T_{1}\) and \(j=y_{i}\) that

\[\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\xi}_{i}\rangle- \langle\mathbf{w}_{j,r}^{(0,0)},\boldsymbol{\xi}_{i}\rangle\] \[\geq\frac{1}{P-1}\rho_{j,r,i}^{(t,b)}-\gamma_{j,r}^{(t,b)}\cdot\| \boldsymbol{\mu}\|_{2}^{-2}\cdot|\langle\boldsymbol{\mu},\boldsymbol{\xi}_{i} \rangle|-\frac{1}{P-1}\sum_{i\neq i^{\prime}}|\rho_{j,r,i^{\prime}}^{(t,b)}| \cdot\|\boldsymbol{\xi}_{i^{\prime}}\|_{2}^{-2}\cdot|\langle\boldsymbol{\xi}_{i ^{\prime}},\boldsymbol{\xi}_{i}\rangle|\] \[\geq-\gamma_{j,r}^{(t,b)}\cdot\|\boldsymbol{\mu}\|_{2}^{-2}\cdot| \langle\boldsymbol{\mu},\boldsymbol{\xi}_{i}\rangle|-\frac{1}{P-1}\sum_{i\neq i ^{\prime}}|\rho_{j,r,i^{\prime}}^{(t,b)}|\cdot\|\boldsymbol{\xi}_{i^{\prime}} \|_{2}^{-2}\cdot|\langle\boldsymbol{\xi}_{i^{\prime}},\boldsymbol{\xi}_{i}\rangle|\] \[\geq-\frac{1}{12}\|\boldsymbol{\mu}\|_{2}^{-2}\cdot|\langle \boldsymbol{\mu},\boldsymbol{\xi}_{i}\rangle|-\frac{n}{12(P-1)}\|\boldsymbol{ \xi}_{i^{\prime}}\|_{2}^{-2}\cdot|\langle\boldsymbol{\xi}_{i^{\prime}}, \boldsymbol{\xi}_{i}\rangle|\] \[\geq-\frac{1}{12}\|\boldsymbol{\mu}\|_{2}^{-1}\sigma_{p}\cdot \sqrt{2\log(6n/\delta)}-\frac{n}{12(P-1)}(\sigma_{p}^{2}d/2)^{-1}2\sigma_{p}^{ 2}\cdot\sqrt{d\log(6n^{2}/\delta)}\] \[=-\frac{1}{12}\|\boldsymbol{\mu}\|_{2}^{-1}\sigma_{p}\cdot\sqrt{2 \log(6n/\delta)}-\frac{n}{3(P-1)}\sqrt{\log(6n^{2}/\delta)/d}\] \[\geq-\frac{1}{6}\|\boldsymbol{\mu}\|_{2}^{-1}\sigma_{p}\cdot\sqrt{ \log(6n/\delta)}.\]

Now let \(\bar{S}_{i}^{(0,0)}\) denote \(\{r:\langle\mathbf{w}_{i,r}^{(0,0)},\boldsymbol{\xi}_{i}\rangle>\sigma_{0} \sigma_{p}\sqrt{d}\}\) and let \(\bar{S}_{j,r}^{(0,0)}\) denote \(\{i\in[n]:y_{i}=j,\;\langle\mathbf{w}_{yi,r}^{(t,b)},\boldsymbol{\xi}_{i} \rangle>\sigma_{0}\sigma_{p}\sqrt{d}\}\). By the condition on \(\sigma_{0}\), we have for \(t\leq T_{1}\) that

\[\langle\mathbf{w}_{j,r}^{(t,b)},\boldsymbol{\xi}_{i}\rangle\geq\frac{1}{\sqrt{ 2}}\langle\mathbf{w}_{j,r}^{(0,0)},\boldsymbol{\xi}_{i}\rangle,\]for any \(r\in\bar{S}_{i}^{(0,0)}\) or \(i\in\bar{S}_{j,r}^{(0,0)}\). Therefore, we have \(\bar{S}_{i}^{(0,0)}\subseteq S_{i}^{(T_{1},0)}\) and \(\bar{S}_{j,r}^{(0,0)}\subseteq S_{j,r}^{(T_{1},0)}\) and hence

\[0.8\Phi(-\sqrt{2})m\leq|\bar{S}_{i}^{(0,0)}|\leq|S_{i}^{(T_{1},0 )}|=\Omega(m),\] \[0.25\Phi(-\sqrt{2})n\leq|\bar{S}_{j,r}^{(0,0)}|\leq|S_{j,r}^{(T_{ 1},0)}|=\Omega(n),\]

where \(\Phi(\cdot)\) is the CDF of the standard normal distribution.

Now we can give proof of Theorem 4.

Proof of Theorem 4.: After the training process of SAM after \(T_{1}\), we get \(\mathbf{W}^{(T_{1},0)}\). To differentiate the SAM process and SGD process. We use \(\widetilde{\mathbf{W}}\) to denote the trajectory obtained by SAM in the proof, i.e., \(\widetilde{\mathbf{W}}^{(T_{1},0)}\). By Proposition D.2, we have that

\[\widetilde{\mathbf{w}}_{j,r}^{(T_{1},0)}=\widetilde{\mathbf{w}}_{j,r}^{(0,0)}+ j\cdot\widetilde{\gamma}_{j,r}^{(T_{1},0)}\cdot\frac{\mathbf{\mu}}{\|\mathbf{\mu}\|_{2}^{2} }+\frac{1}{P-1}\sum_{i=1}^{n}\widetilde{\rho}_{j,r,i}^{(T_{1},0)}\cdot\frac{\bm {\xi}_{i}}{\|\mathbf{\xi}_{i}\|_{2}^{2}}+\frac{1}{P-1}\sum_{i=1}^{n}\widetilde{\rho }_{j,r,i}^{(T_{1},0)}\cdot\frac{\mathbf{\xi}_{i}}{\|\mathbf{\xi}_{i}\|_{2}^{2}}, \tag{69}\]

where \(\widetilde{\gamma}_{j,r}^{(T_{1},0)}=\Theta(1)\), \(\widetilde{\rho}_{j,r,i}^{(T_{1},0)}\in[0,1/12]\), \(\widetilde{\rho}_{j,r,i}^{(T_{1},0)}\in[-\beta-10\sqrt{\log(6n^{2}/\delta)/d}n,0]\). Then the SGD start at \(\mathbf{W}^{(0,0)}:=\widetilde{\mathbf{W}}^{(T_{1},0)}\). Notice that by Lemma D.7, we know that the initial weights of SGD (i.e., the end weight of SAM) \(\mathbf{W}^{(0,0)}\) still satisfies the conditions for Subsection C.1 and C.2. Therefore, following the same analysis in Subsection C.1 and C.2, we have that there exist \(t=\widetilde{O}(\eta^{-1}\epsilon^{-1}mnd^{-1}P^{-2}\sigma_{p}^{-2})\) such that \(L_{S}(\mathbf{W}^{(t,0)})\leq\epsilon\). Besides,

\[\mathbf{w}_{j,r}^{(t,0)}=\mathbf{w}_{j,r}^{(0,0)}+j\cdot\gamma_{j,r}^{(t,0)} \cdot\frac{\mathbf{\mu}}{\|\mathbf{\mu}\|_{2}^{2}}+\frac{1}{P-1}\sum_{i=1}^{n}\widetilde {\rho}_{j,r,i}^{(t,0)}\cdot\frac{\mathbf{\xi}_{i}}{\|\mathbf{\xi}_{i}\|_{2}^{2}}+\frac{ 1}{P-1}\sum_{i=1}^{n}\rho_{j,r,i}^{(t,0)}\cdot\frac{\mathbf{\xi}_{i}}{\|\mathbf{\xi}_{ i}\|_{2}^{2}}\]

for \(j\in[\pm 1]\) and \(r\in[m]\) where

\[\gamma_{j,r}^{(t,0)}=\Theta(\mathrm{SNR}^{2})\sum_{i\in[n]}\overline{\rho}_{j,r,i}^{(t,0)},\quad\overline{\rho}_{j,r,i}^{(t,0)}\in[0,\alpha],\quad\underline {\rho}_{j,r,i}^{(t,0)}\in[-\alpha,0]. \tag{70}\]

Next, we will evaluate the test error for \(\mathbf{W}^{(t,0)}\). Notice that we use \((t)\) as the shorthand notation of \((t,0)\). For the sake of convenience, we use \((\mathbf{x},\widehat{y},y)\sim\mathcal{D}\) to denote the following: data point \((\mathbf{x},y)\) follows distribution \(\mathcal{D}\) defined in Definition 2.1, and \(\widehat{y}\) is its true label. We can write out the test error as

\[\mathbb{P}_{(\mathbf{x},y)\sim\mathcal{D}}\big{(}y\neq\operatorname {sign}(f(\mathbf{W}^{(t)},\mathbf{x}))\big{)}\] \[=\mathbb{P}_{(\mathbf{x},y)\sim\mathcal{D}}\big{(}yf(\mathbf{W}^{ (t)},\mathbf{x})\leq 0\big{)}\] \[=\mathbb{P}_{(\mathbf{x},y)\sim\mathcal{D}}\big{(}yf(\mathbf{W}^{ (t)},\mathbf{x})\leq 0,y\neq\widehat{y}\big{)}+\mathbb{P}_{(\mathbf{x},\widehat{y},y )\sim\mathcal{D}}\big{(}yf(\mathbf{W}^{(t)},\mathbf{x})\leq 0,y=\widehat{y}\big{)} \tag{71}\] \[=p\cdot\mathbb{P}_{(\mathbf{x},\widehat{y},y)\sim\mathcal{D}} \big{(}\widehat{y}f(\mathbf{W}^{(t)},\mathbf{x})\geq 0\big{)}+(1-p)\cdot \mathbb{P}_{(\mathbf{x},\widehat{y},y)\sim\mathcal{D}}\big{(}\widehat{y}f( \mathbf{W}^{(t)},\mathbf{x})\leq 0\big{)}\] \[\leq p+\mathbb{P}_{(\mathbf{x},\widehat{y},y)\sim\mathcal{D}} \big{(}\widehat{y}f(\mathbf{W}^{(t)},\mathbf{x})\leq 0\big{)},\]

where in the second equation we used the definition of \(\mathcal{D}\) in Definition 2.1. It therefore suffices to provide an upper bound for \(\mathbb{P}_{(\mathbf{x},\widehat{y})\sim\mathcal{D}}\big{(}\widehat{y}f( \mathbf{W}^{(t)},\mathbf{x})\leq 0\big{)}\). To achieve this, we write \(\mathbf{x}=(\widehat{y}\mathbf{\mu},\mathbf{\xi})\), and get

\[\widehat{y}f(\mathbf{W}^{(t)},\mathbf{x}) =\frac{1}{m}\sum_{j,r}\widehat{y}j[\sigma(\langle\mathbf{w}_{j, r}^{(t)},\widehat{y}\mathbf{\mu}\rangle)+\sigma(\langle\mathbf{w}_{j,r}^{(t)},\mathbf{\xi} \rangle)]\] \[=\frac{1}{m}\sum_{r}[\sigma(\langle\mathbf{w}_{\widehat{y},r}^{(t )},\widehat{y}\mathbf{\mu}\rangle)+(P-1)\sigma(\langle\mathbf{w}_{\widehat{y},r}^{(t )},\mathbf{\xi}\rangle)]\] \[\qquad-\frac{1}{m}\sum_{r}[\sigma(\langle\mathbf{w}_{-\widehat{y},r}^{(t)},\widehat{y}\mathbf{\mu}\rangle)+(P-1)\sigma(\langle\mathbf{w}_{-\widehat{y},r}^{(t)},\mathbf{\xi}\rangle)]. \tag{72}\]The inner product with \(j=\widehat{y}\) can be bounded as

\[\langle\mathbf{w}^{(t)}_{\widehat{y},r},\widehat{y}\boldsymbol{\mu}\rangle =\langle\mathbf{w}^{(0)}_{\widehat{y},r},\widehat{y}\boldsymbol{ \mu}\rangle+\gamma^{(t)}_{\widehat{y},r}+\frac{1}{(P-1)}\sum_{i=1}^{n}\overline {\rho}^{(t)}_{\widehat{y},r,i}\cdot\|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot \langle\boldsymbol{\xi}_{i},\widehat{y}\boldsymbol{\mu}\rangle \tag{73}\] \[\qquad+\frac{1}{(P-1)}\sum_{i=1}^{n}\underline{\rho}^{(t)}_{ \widehat{y},r,i}\cdot\|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot\langle\boldsymbol{ \xi}_{i},\widehat{y}\boldsymbol{\mu}\rangle\] \[\geq\langle\mathbf{w}^{(0)}_{\widehat{y},r},\widehat{y} \boldsymbol{\mu}\rangle+\gamma^{(t)}_{\widehat{y},r}-\frac{\sqrt{2\log(6n/ \delta)}}{P-1}\cdot\sigma_{p}\|\boldsymbol{\mu}\|_{2}\cdot(\sigma_{p}^{2}d/2 )^{-1}\bigg{[}\sum_{i=1}^{n}\overline{\rho}^{(t)}_{\widehat{y},r,i}+\sum_{i=1 }^{n}|\underline{\rho}^{(t)}_{\widehat{y},r,i}|\bigg{]}\] \[=\langle\mathbf{w}^{(0)}_{\widehat{y},r},\widehat{y}\boldsymbol{ \mu}\rangle+\gamma^{(t)}_{\widehat{y},r}-\Theta\big{(}\sqrt{\log(n/\delta)} \cdot(P\sigma_{p}d)^{-1}\|\boldsymbol{\mu}\|_{2}\big{)}\cdot\Theta(\mathrm{SNR }^{-2})\cdot\gamma^{(t)}_{\widehat{y},r}\] \[=\langle\mathbf{w}^{(0)}_{\widehat{y},r},\widehat{y}\boldsymbol{ \mu}\rangle+\big{[}1-\Theta\big{(}\sqrt{\log(n/\delta)}\cdot P\sigma_{p}/\| \boldsymbol{\mu}\|_{2}\big{)}\big{]}\gamma^{(t)}_{\widehat{y},r}\] \[=\langle\mathbf{w}^{(0)}_{\widehat{y},r},\widehat{y}\boldsymbol{ \mu}\rangle+\Theta(\gamma^{(t)}_{\widehat{y},r})\] \[=\Omega(1),\]

where the inequality is by Lemma B.1; the second equality is obtained by plugging in the coefficient orders we summarized at (70); the third equality is by the condition \(\mathrm{SNR}=\|\boldsymbol{\mu}\|_{2}/P\sigma_{p}\sqrt{d}\); the fourth equality is due to \(\|\boldsymbol{\mu}\|_{2}^{2}\geq C\cdot P^{2}\sigma_{p}^{2}\log(n/\delta)\) in Condition 3.1, so for sufficiently large constant \(C\) the equality holds; the last equality is by Lemma D.7. Moreover, we can deduce in a similar manner that

\[\langle\mathbf{w}^{(t)}_{-\widehat{y},r},\widehat{y}\boldsymbol{ \mu}\rangle =\langle\mathbf{w}^{(0)}_{-\widehat{y},r},\widehat{y} \boldsymbol{\mu}\rangle-\gamma^{(t)}_{-\widehat{y},r}+\sum_{i=1}^{n}\overline {\rho}^{(t)}_{-\widehat{y},r,i}\cdot\|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot \langle\boldsymbol{\xi}_{i},-\widehat{y}\boldsymbol{\mu}\rangle \tag{74}\] \[\qquad+\sum_{i=1}^{n}\underline{\rho}^{(t)}_{-\widehat{y},r,i} \cdot\|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot\langle\boldsymbol{\xi}_{i}, \widehat{y}\boldsymbol{\mu}\rangle\] \[\leq\langle\mathbf{w}^{(0)}_{-\widehat{y},r},\widehat{y} \boldsymbol{\mu}\rangle-\gamma^{(t)}_{-\widehat{y},r}\] \[\qquad+\sqrt{2\log(6n/\delta)}\cdot\sigma_{p}\|\boldsymbol{\mu} \|_{2}\cdot(\sigma_{p}^{2}d/2)^{-1}\bigg{[}\sum_{i=1}^{n}\overline{\rho}^{(t )}_{-\widehat{y},r,i}+\sum_{i=1}^{n}|\underline{\rho}^{(t)}_{-\widehat{y},r, i}|\bigg{]}\] \[=\langle\mathbf{w}^{(0)}_{-\widehat{y},r},\widehat{y}\boldsymbol {\mu}\rangle-\Theta(\gamma^{(t)}_{-\widehat{y},r})\] \[=-\Omega(1)<0,\]

where the second equality holds based on similar analyses as in (73).

Denote \(g(\boldsymbol{\xi})\) as \(\sum_{r}\sigma(\langle\mathbf{w}^{(t)}_{-\widehat{y},r},\boldsymbol{\xi} \rangle)\). According to Theorem 5.2.2 in Vershynin (2018), we know that for any \(x\geq 0\) it holds that

\[\mathbb{P}(g(\boldsymbol{\xi})-\mathbb{E}g(\boldsymbol{\xi})\geq x)\leq\exp \Big{(}-\frac{cx^{2}}{\sigma_{p}^{2}\|g\|_{\mathrm{Lip}}^{2}}\Big{)}, \tag{75}\]

where \(c\) is a constant. To calculate the Lipschitz norm, we have

\[|g(\boldsymbol{\xi})-g(\boldsymbol{\xi}^{\prime})| =\bigg{|}\sum_{r=1}^{m}\sigma(\langle\mathbf{w}^{(t)}_{-\widehat {y},r},\boldsymbol{\xi}\rangle)-\sum_{r=1}^{m}\sigma(\langle\mathbf{w}^{(t)} _{-\widehat{y},r},\boldsymbol{\xi}^{\prime}\rangle)\bigg{|}\] \[\leq\sum_{r=1}^{m}\big{|}\sigma(\langle\mathbf{w}^{(t)}_{- \widehat{y},r},\boldsymbol{\xi}\rangle)-\sigma(\langle\mathbf{w}^{(t)}_{- \widehat{y},r},\boldsymbol{\xi}^{\prime}\rangle)\big{|}\] \[\leq\sum_{r=1}^{m}|\langle\mathbf{w}^{(t)}_{-\widehat{y},r}, \boldsymbol{\xi}-\boldsymbol{\xi}^{\prime}\rangle|\] \[\leq\sum_{r=1}^{m}\big{\|}\mathbf{w}^{(t)}_{-\widehat{y},r}\|_{2 }\cdot\|\boldsymbol{\xi}-\boldsymbol{\xi}^{\prime}\|_{2},\]where the first inequality is by triangle inequality, the second inequality is by the property of ReLU; and the last inequality is by Cauchy-Schwartz inequality. Therefore, we have

\[\|g\|_{\mathrm{Lip}}\leq\sum_{r=1}^{m}\big{\|}\mathbf{w}_{-\widehat{y},r}^{(t)} \big{\|}_{2}, \tag{76}\]

and since \(\langle\mathbf{w}_{-\widehat{y},r}^{(t)},\boldsymbol{\xi}\rangle\sim\mathcal{N }\big{(}0,\|\mathbf{w}_{-\widehat{y},r}^{(t)}\|_{2}^{2}\sigma_{p}^{2}\big{)}\), we can get

\[\mathbb{E}g(\boldsymbol{\xi})=\sum_{r=1}^{m}\mathbb{E}\sigma(\langle\mathbf{w}_{ -\widehat{y},r}^{(t)},\boldsymbol{\xi}\rangle)=\sum_{r=1}^{m}\frac{\|\mathbf{ w}_{-\widehat{y},r}^{(t)}\|_{2}\sigma_{p}}{\sqrt{2\pi}}=\frac{\sigma_{p}}{\sqrt{2 \pi}}\sum_{r=1}^{m}\|\mathbf{w}_{-\widehat{y},r}^{(t)}\|_{2}.\]

Next we seek to upper bound the \(2\)-norm of \(\mathbf{w}_{j,r}^{(t)}\). First, we tackle the noise section in the decomposition, namely:

\[\bigg{\|}\sum_{i=1}^{n}\rho_{j,r,i}^{(t)}\cdot\|\boldsymbol{\xi} \|_{2}^{-2}\cdot\boldsymbol{\xi}_{i}\bigg{\|}_{2}^{2}\] \[=\sum_{i=1}^{n}{\rho_{j,r,i}^{(t)}}^{2}\cdot\|\boldsymbol{\xi}_{i }\|_{2}^{-2}+2\sum_{1\leq i_{1}<i_{2}\leq n}\rho_{j,r,i_{1}}^{(t)}\rho_{j,r,i_ {2}}^{(t)}\cdot\|\boldsymbol{\xi}_{i_{1}}\|_{2}^{-2}\cdot\|\boldsymbol{\xi}_{i _{2}}\|_{2}^{-2}\cdot\langle\boldsymbol{\xi}_{i_{1}},\boldsymbol{\xi}_{i_{2}}\rangle\] \[\leq 4\sigma_{p}^{-2}d^{-1}\sum_{i=1}^{n}{\rho_{j,r,i}^{(t)}}^{2 }+2\sum_{1\leq i_{1}<i_{2}\leq n}\big{|}\rho_{j,r,i_{1}}^{(t)}\rho_{j,r,i_{2}} ^{(t)}\big{|}\cdot(16\sigma_{p}^{-4}d^{-2})\cdot(2\sigma_{p}^{2}\sqrt{d\log(6 n^{2}/\delta)})\] \[=4\sigma_{p}^{-2}d^{-1}\sum_{i=1}^{n}{\rho_{j,r,i}^{(t)}}^{2}+32 \sigma_{p}^{-2}d^{-3/2}\sqrt{\log(6n^{2}/\delta)}\bigg{[}\bigg{(}\sum_{i=1}^{n }\big{|}\rho_{j,r,i}^{(t)}\big{|}\bigg{)}^{2}-\sum_{i=1}^{n}{\rho_{j,r,i}^{(t )}}^{2}\bigg{]}\] \[=\Theta(\sigma_{p}^{-2}d^{-1})\sum_{i=1}^{n}{\rho_{j,r,i}^{(t)}}^ {2}+\widetilde{\Theta}(\sigma_{p}^{-2}d^{-3/2})\bigg{(}\sum_{i=1}^{n}\big{|} \rho_{j,r,i}^{(t)}\bigg{)}^{2}\] \[\leq\Theta(\sigma_{p}^{-2}d^{-1}n^{-1})\bigg{(}\sum_{i=1}^{n} \overline{\rho}_{j,r,i}^{(t)}\bigg{)}^{2},\]

where for the first inequality, we used Lemma B.1; for the second inequality, we used the definition of \(\overline{\rho}\), \(\underline{\rho}\); for the second to last equation, we plugged in coefficient orders. We can thus upper bound the \(2\)-norm of \(\mathbf{w}_{j,r}^{(t)}\) as:

\[\|\mathbf{w}_{j,r}^{(t)}\|_{2} \leq\|\mathbf{w}_{j,r}^{(0)}\|_{2}+\gamma_{j,r}^{(t)}\cdot\| \boldsymbol{\mu}\|_{2}^{-1}+\frac{1}{P-1}\bigg{\|}\sum_{i=1}^{n}\rho_{j,r,i}^ {(t)}\cdot\|\boldsymbol{\xi}_{i}\|_{2}^{-2}\cdot\boldsymbol{\xi}_{i}\bigg{\|} _{2}\] \[\leq\|\mathbf{w}_{j,r}^{(0)}\|_{2}+\gamma_{j,r}^{(t)}\cdot\| \boldsymbol{\mu}\|_{2}^{-1}+\Theta(P^{-1}\sigma_{p}^{-1}d^{-1/2}n^{-1/2})\cdot \sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(t)}\] \[=\Theta(\sigma_{0}\sqrt{d})+\Theta(P^{-1}\sigma_{p}^{-1}d^{-1/2} n^{-1/2})\cdot\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(t)}, \tag{77}\]

where the first inequality is due to the triangle inequality, and the equality is due to the following comparisons:

\[\frac{\gamma_{j,r}^{(t)}\cdot\|\boldsymbol{\mu}\|_{2}^{-1}}{ \Theta(P^{-1}\sigma_{p}^{-1}d^{-1/2}n^{-1/2})\cdot\sum_{i=1}^{n}\overline{\rho }_{j,r,i}^{(t)}} =\Theta(P^{-1}\sigma_{p}d^{1/2}n^{1/2}\|\boldsymbol{\mu}\|_{2}^{ -1}\mathrm{SNR}^{2})\] \[=\Theta(P^{-1}\sigma_{p}^{-1}d^{-1/2}n^{1/2}\|\boldsymbol{\mu}\|_ {2})\] \[=O(1)\]based on the coefficient order \(\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{(t)}/\gamma_{j,r}^{(t)}=\Theta(\mathrm{SNR}^ {-2})\), the definition \(\mathrm{SNR}=\|\mathbf{\mu}\|_{2}/(\sigma_{p}\sqrt{d})\), and the condition for \(d\) in Condition 3.1; and also \(\|\mathbf{w}_{j,r}^{(0)}\|_{2}=\Theta(\sigma_{0}\sqrt{d})\) based on Lemma D.7. With this and (73), we analyze the key component in (80):

\[\frac{\sum_{r}\sigma(\langle\mathbf{w}_{\widehat{y},r}^{(t)}, \widehat{y}\mathbf{\mu}\rangle)}{(P-1)\sigma_{p}\sum_{r=1}^{m}\left\|\mathbf{w}_{- \widehat{y},r}^{(t)}\right\|_{2}} \geq\frac{\Theta(1)}{\Theta(\sigma_{0}\sqrt{d})+\Theta(P^{-1} \sigma_{p}^{-1}d^{-1/2}n^{-1/2})\cdot\sum_{i=1}^{n}\overline{\rho}_{j,r,i}^{( t)}}\] \[\geq\frac{\Theta(1)}{\Theta(\sigma_{0}\sqrt{d})+O(P^{-1}\sigma_{p }^{-1}d^{-1/2}n^{1/2}\alpha)}\] \[\geq\min\{\Omega(\sigma_{0}^{-1}d^{-1/2}),\Omega(P\sigma_{p}d^{1 /2}n^{-1/2}\alpha^{-1})\}\] \[\geq 1.\]

It directly follows that

\[\sum_{r}\sigma(\langle\mathbf{w}_{\widehat{y},r}^{(t)},\widehat{y}\mathbf{\mu} \rangle)-\frac{(P-1)\sigma_{p}}{\sqrt{2\pi}}\sum_{r=1}^{m}\|\mathbf{w}_{- \widehat{y},r}^{(t)}\|_{2}>0. \tag{79}\]

Now using the method in (75) with the results above, we plug (74) into (72) and then (71), to obtain

\[\mathbb{P}_{(\mathbf{x},\widehat{y},y)\sim\mathcal{D}}\Big{(} \widehat{y}f(\mathbf{W}^{(t)},\mathbf{x})\leq 0\Big{)}\] \[=\mathbb{P}_{(\mathbf{x},\widehat{y},y)\sim\mathcal{D}}\Bigg{(}g( \mathbf{\xi})-\mathbb{E}g(\mathbf{\xi})\geq(1/(P-1))\sum_{r}\sigma(\langle\mathbf{w}_{ \widehat{y},r}^{(t)},\widehat{y}\mathbf{\mu}\rangle)-\frac{\sigma_{p}}{\sqrt{2\pi }}\sum_{r=1}^{m}\|\mathbf{w}_{-\widehat{y},r}^{(t)}\|_{2}\Bigg{)}\] \[\leq\exp\Bigg{[}-\frac{c\Big{(}(1/(P-1))\sum_{r}\sigma(\langle \mathbf{w}_{\widehat{y},r}^{(t)},\widehat{y}\mathbf{\mu}\rangle)-(\sigma_{p}/ \sqrt{2\pi})\sum_{r=1}^{m}\left\|\mathbf{w}_{-\widehat{y},r}^{(t)}\right\|_{2 }\Big{)}^{2}}{\sigma_{p}^{2}\Big{(}\sum_{r=1}^{m}\left\|\mathbf{w}_{-\widehat{ y},r}^{(t)}\right\|_{2}\Big{)}^{2}}\Bigg{]}\] \[=\exp\bigg{[}-c\bigg{(}\frac{\sum_{r}\sigma(\langle\mathbf{w}_{ \widehat{y},r}^{(t)},\widehat{y}\mathbf{\mu}\rangle)}{(P-1)\sigma_{p}\sum_{r=1}^{ m}\left\|\mathbf{w}_{-\widehat{y},r}^{(t)}\right\|_{2}}-1/\sqrt{2\pi}\bigg{)}^{2} \bigg{]}\] \[\leq\exp(c/2\pi)\exp\bigg{(}-0.5c\bigg{(}\frac{\sum_{r}\sigma( \langle\mathbf{w}_{\widehat{y},r}^{(t)},\widehat{y}\mathbf{\mu}\rangle)}{(P-1) \sigma_{p}\sum_{r=1}^{m}\left\|\mathbf{w}_{-\widehat{y},r}^{(t)}\right\|_{2} }\bigg{)}^{2}\bigg{)}, \tag{80}\]

where the second inequality is by (79) and plugging (76) into (75), the third inequality is due to the fact that \((s-t)^{2}\geq s^{2}/2-t^{2},\forall s,t\geq 0\).

And we can get from (78) and (80) that

\[\mathbb{P}_{(\mathbf{x},\widehat{y},y)\sim\mathcal{D}}\big{(} \widehat{y}f(\mathbf{W}^{(t)},\mathbf{x})\leq 0\big{)} \leq\exp(c/2\pi)\exp\bigg{(}-0.5c\bigg{(}\frac{\sum_{r}\sigma( \langle\mathbf{w}_{\widehat{y},r}^{(t)},\widehat{y}\mathbf{\mu}\rangle)}{(P-1) \sigma_{p}\sum_{r=1}^{m}\left\|\mathbf{w}_{-\widehat{y},r}^{(t)}\right\|_{2}} \bigg{)}^{2}\bigg{)}\] \[\leq\exp\Big{(}\frac{c}{2\pi}-C\min\{\sigma_{0}^{-2}d^{-1},P \sigma_{p}^{2}dn^{-1}\alpha^{-2}\}\Big{)}\] \[\leq\exp\Big{(}-0.5C\min\{\sigma_{0}^{-2}d^{-1},P\sigma_{p}^{2}dn ^{-1}\alpha^{-2}\}\Big{)}\] \[\leq\epsilon,\]

where \(C=O(1)\), the last inequality holds since \(\sigma_{0}^{2}\leq 0.5Cd^{-1}\log(1/\epsilon)\) and \(d\geq 2C^{-1}P^{-1}\sigma_{p}^{-2}n\alpha^{2}\log(1/\epsilon)\).

Response:
## Review

### Summary

This paper presents a theoretical examination of Sharpness-Aware Minimization (SAM) in the context of two-layer convolutional ReLU networks, contrasting its effectiveness against Stochastic Gradient Descent (SGD). It provides conditions under which benign overfitting occurs with both methods, demonstrating that SAM can achieve better generalization, particularly in scenarios where SGD leads to harmful overfitting. The authors substantiate their claims with numerical experiments, revealing SAM's advantages in mitigating noise learning and enhancing the learning of weak features. Though the contributions are significant and well-structured, the presentation requires improvement to enhance clarity and definition consistency.

### Strengths

- The paper addresses the significant issue of overfitting in large neural networks, making it timely and relevant.
- It provides a strong theoretical analysis explaining why SAM outperforms SGD, enhancing understanding of generalization in neural networks.
- The authors conduct a comprehensive study comparing SAM and SGD using both synthetic and real data, adding validity to their findings.
- The work claims to present the first benign overfitting results for mini-batch SGD, which could signify a novel contribution to the field.
- The paper is well-structured and clearly written, providing a good overview of background literature and detailed contributions.

### Weaknesses

- The presentation is disordered and lacks clear definitions for several notations, making the results difficult to understand.
- Certain key theoretical aspects, such as the harmful overfitting regime for SAM, are not sufficiently covered.
- The experiments may lack reliability due to potentially high learning rates, which could misrepresent the findings.
- The paper's focus on a specific architecture limits the generalizability of its findings to other neural network types.
- The related work section is more of an enumeration rather than a discussion of technical differences with prior works.

### Questions

- Why are equations in Line 175 and 176 presented as definitions when they seem to stem from the data distribution and network architecture?
- Could the authors clarify why the learning rate was set to 0.1 in the experiments, given concerns over its potential to cause overfitting?
- Is the theoretical result for SGD also applicable to Gradient Flow dynamics?
- Will switching from SAM to SGD after a few epochs yield similar performance to SAM?
- How does the perturbation radius (τ) requirement in theory reconcile with practical applications?
- Could the authors elaborate on the connection between SAM and recent work regarding large learning rates and their implications for generalization?

### Soundness

**Score:** 3

**Description:** 3 = good; the theoretical analysis is solid, but some parts lack completeness and robustness.

### Presentation

**Score:** 3

**Description:** 3 = good; while the paper is well-structured, clarity could be improved and some notations need definition.

### Contribution

**Score:** 4

**Description:** 4 = excellent; the paper makes significant contributions to understanding SAM's advantages over SGD, particularly in benign overfitting.

### Rating

**Score:** 7

**Description:** 7 = accept; the paper is technically solid with a high impact on the field, but requires minor improvements in presentation and evaluation.

### Paper Decision

**Decision:** Accept (poster)

**Reasons:** The paper presents important theoretical contributions to understanding SAM's advantages over SGD in mitigating overfitting, supported by numerical experiments. While the results are significant, the presentation requires clarity and definition consistency. Given the solid theoretical grounding and the potential impact on the field, the decision is to accept the paper with suggestions for improving presentation and addressing the raised questions.