Why Does Sharpness-Aware Minimization Generalize Better Than SGD?
Authors: Zixiang Chen
Junkai Zhang
Yiwen Kou
Xiangning Chen
Cho-Jui Hsieh
Quanquan Gu
Venue: NeurIPS 2023 poster (2023)
Keywords: Sharpness Aware Algorithm
Deep Learning Theory

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.

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, mannzx,,,,hann,,.. evankou, xiangning, chohsieh, qgu>ecs. wcw..claee edu
Abstract
The challenge 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 promising training method, which can improve the generalization of neural networks even in the presence of label noise. However, 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 theory.
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 unstable and difficult due to the sharpness, non-convexity, and non-smoothness of its loss landscape. In addition, the number of model parameters is much larger than the training sample size, the model has the ability to memorize randomly labeled data (Zhang al.) 2021), which leads to overfitting. Therefore, although traditional gradient-basse 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, Charpness-Awar Minimization (SAM) (Foret et al. 2020 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 solution with reduced sharpness, characterized by how rapidly the loss function changes locally. Despite the empirical success SAM across numerous tasks (Bahri et al., 2021] [ehdin al. 2022 Chen al. limited.
2021 Liu al. 2022a), the theoretical understanding of this method remains
Foret et al. (2020) provided a PAC-Bayes bound on the generalization error of SAM show that it will generalize well, while the bound only holds for the infeasible average-direction perturbation instead of
'Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

SGD
SAM
Harmful Overfitting
|| 201)
Dimension
Dimenalon d
|| O(d
Benign Overfitting
Benign Overfitting
Signal Strength lprll
Signal Strength ||
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 guaranteed to have 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.
practically used ascend-direction perturbation. Andriushchenko and Flammarion (002) 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 theoretical basis demonstrating when SAM outperforms SGD. In particular, we consider a data distribution mainly characterized by the signal ls 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.1p. Let p be the strength of the label flipping noise. For any e > O, under certain regularity conditions, with high probability, there exists 0 S t T such that the training loss converges, i.e., Ls(W()) I Besides,
1. For SGD, when the signal strength ||al (dd1/1), have T.g'((W(0), S Sptc t... When the signal strength ||l||2 Od)/4), we have D3 (W(1) 2 2 0.1.
For SAM, provided the signal strength || 2 (11), we have Lb-1 (W(1) sptc
Our contributions summarized 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 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-bater batch SGD. We also prove a phase transition phenomenon for SGD, which illustrated in Figure] I
Under the conditions when SGD leads to harmful overfitting, we formally prove that SAM can achieve benign overfitting. Consequently, establish rigorous theoretical distinction between SAM and SGD, demonstrating that SAM strictly outperforms SGD in terms of generalization error. Specifically, 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 vector (V).--- va) we denote

---

by livl2 == (Σ)= 03) its 2-norm. For two sequence {a*} and {6x J. we denote as (bb)) if asl < Clbba for some absolute constant C, denote a = (bbb) if by O(ay), and denote as (bb)) if az O(b)) and Oa 9(bx). We also denote au - o(ba) lim lax/bal = O. Finally, use (()) and ((-) to omit logarithmic terms in the notation. We denote the set (1,..., n) by [n]. and the set {0. ....... n 1} by [n], respectively. The carnality of set S is denoted by |S].
Preliminaries
2.1 Data distribution
Our focus is on binary classification with label y E {+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 μ E Rd be a fixed vector representing the signal contained in each data point. Each
data point (x. y) with input x [x(1)T xx2)T TP))T] E Rxd x(1), x(2),
, x(P) E R
and label y E 1 is generated from a distribution D specified as follows:
1. The true label ÿ is generated as Rademacher random variable, i.e., Ply 1] = PIù - 1] = 1/2. The observed label is then generated by flipping 7 with probability where p < 1/2, i.e., Ply = 5 = 1-11-- p and IPly - - p.
2. A noise vector ξ is generated from the Gaussian distribution (N o2), where of is the variance. is randomly selected and then assigned as y j4, which represents the signal, while the others are given by &, which represents noises.
3. One of x(),, x(2) ...... x(
The data distribution in Definition 2.1 I 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. (023).
When P 2, this data distribution aligns with the analyzed in Kou et al. (023) 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 x = (x())._ ..... x(P) there is exactly one x(i) y/4, and the others are random Gaussian vectors. For example, the input vector x can be or [E, ........ &, yul: the signal patch you can appear at any position. To avoid harmful overfitting, the
El..... [..... yuu, E]
model must learn the feature vector rather than the noise vector.
2.2 Neural Network and Training Loss
To effectively learn the distribution per Definition 2.1 it advantageous to utilize shared weights structure, given that the specific signal patch is not known beforehand. When P >> n, shared weights become indispensable the location of the signal patch in the test can differ from the location of the signal patch in the training data.
We consider two-layer convolutional neural network whose filters are applied to the P patches X1, ... X2 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 x) and F-1(W (W-,, x) are defined
f(W, x) = FF1((W (1. - -F F-1(W- (W-1. where F41 (W
W ,) = m-mmmiaa -E- E P
p= ((W x())))
Here we consider ReLU activation function ((2) 1(: 2 0)z, Wjr Rd denotes the weight for the r-th filter, and W, is the collection of model weights associated with F, for j = #1. We use W to denote the collection of all model weights. Denote the training data set by S ((x,,)) [e])) We train the above CNN model by minimizing the empirical cross-entropy loss function
Ls(W) EEeen (W,,
where {(z) = log(1 + exp(-2)).
2.3 Training Algorithm
Minibatch Stochastic Gradient Descent. For epoch 1, the training data set randomly divided into H : n/ B mini batches It,b with batch size B > 2. The empirical loss for batch Ie, is defined
3

---

as LI (W) ( (1/B) EiEIe (Lu.f(Wxii)) Then the gradient descent update of the filters in the CNN can be written
w (((++1) w(f.b) - L (W
(2)
where the gradient of the empirical loss LW.., is the collection of V Ww,,- LI,, follows
Wwy LI (
(P-1)
Σ ((ai)) o' d'((w),
(8,0) .Ei) june
Bm
iEIea
(1.b) ((bb) o' d'(W""r (b.l)) TUU5..
(3)
Bm
for all {+1} and r [m]. Here introduce shorthand notation (( l' f( f(W((.)) x,) and assume the gradient of the ReLU activation function at 0 to be o'(0) = 1 without loss of generality. We use (t, b) to denote epoch index f with mini-batch index and use (t) as the shorthand of (t, 0). We initialize SGD by random Gaussian, where all entries of (())
sampled from i.i.d.
Gaussian distributions (00, o)) with oî being the variance. From B we can infer that the loss landscape of the empirical loss highly non-smooth because the ReL.U function is not differentiable at zero. In particular, when ( (f.)))
.5) close to zero, even a very small perturbation can greatly
j.r
change the activation pattern o' d'(ww (t.b) {)) and thus change the direction of Vw,,, V (W(+.)).. j.r
This observation prevents the analysis technique based 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 (W) with trainable parameter W, the idea of SAM ball of W to ensure uniformly low training loss value. In particular, it aims to solve the following optimization problem
minimize a perturbed empirical loss at the worst point in the neighborhood
min W
SAM (W). where S
SAM (W) == max Ls W + c),
|elllaST
where the hyperparameter 7 called the perturbation radius. However, directly optimizing L (W) computationally expensive. In practice, people the following sharpness-aware minimization (SAM) algorithm (Foret et al. 2020 Zheng et al. 2021) to minimize L SAA (W) efficiently,
V (W)
((++)) (()) Ls W + + where =
(5)
V (w
When applied to SGD in 2 the gradient Vw s in (5) is further replaced by stochastic gradient
Ww (Foret al. 2020 The detailed algorithm description of SAM in shown in Algorithm
Algorithm Minibatch Sharpness Aware Minimization
Input: Training set S -
(x,, y:)}. Batch size B, step size 17 0, neighborhood size
0.
Initialize weights (0)) for t = 0, 1,...,
T ,
do
Randomly divide the training data set into mini batches (I for = 0, 1, H do
1-11
Fwhr, (W (Wit.)))
We calculate the perturbation ((t,)) 7 TVWLT, (W
Update model parameters: W(tb+1)) W(i. -1 VV LWL,,, (W) www.0004.0..
end for
Update model parameters: W(++1.0) W(LH)
end for
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 og and learning rate 17.

---

Condition 3.1. Suppose there exists a sufficiently large constant C, such that the following hold:
1. Dimension d sufficiently large: d 2 O max(nP o0, 20,,,11111111111111111 211 n² Po.2,.P-0000 oopp (m,BB)])
2. Training sample size n and neural network width satisfy m.n 9((1)
3. The 2-norm of the signal satisties || 2 (Pppp))
4. The noise rate p satisfies S 1/C.
The standard deviation of Gaussian initialization do appropriately chosen such that do S
ö (max (Pa,d/ vn, l1la))))))
The learning rate 7) satisfies 7) S Ō(( ((max {P2 (Bm))//BBBBm)))) P3o.d/ B, B.lll////
n nPopllul22/(H)) (Pop)llllla((Bmmm)))
The conditions imposed on the data dimensions d, network width m, and the number of samples ensure adequate overparameterization
the network. Additionally, the condition on the learning rate 17 facilitates efficient learning by our model. By concentration inequality, with high probability, the lz norm of the noise patch is of order (do Therefore, the quantity d can be viewed as the strength of the noise. Comparable conditions have been established in Chatterji and Long (202I); Cao al. 20022); Frei al, (2022); Kou et al. (2023). Based on the above condition, we first present 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 > O, under Condi- tion 3.1 with probability at least 1 - there exists t = Ö(1 -mm' ==pp²) such that:
1. The training loss converges, i.e., Ls Ls(W S.
2. When n |1 2 2 C CidPto
o the test error L% (W(") pt
3. When m1ll S C3dPa of. the test error (w(() 2 0.1.
Theorem 3.2 reveals 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 v 11 Z Cdd(P 1) the model can efficiently learn the signal. As 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 f I <. Csd(P= P 1)a of holds, the test error fails to approach the Bayes risk. This phase transition is empirically illustrated in Figure 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- sigh-dimensional data. 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, we will discuss later in Section
3.1 Analysis of Mini-Batel 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:
5

---

Lemma 3.3. Let whr for E (+1), E [m]] be the convolution filters of the CNN at the b-th batch
(1.b)
t- epoch of gradient descent. Then there exist unique coefficients Yisr and Pjrz such that
(t.b) (1,b)
(t.b)
(0,0)
(t.b)
1|121122
+ut P-1
Parg 116/112
(1.b)
(6)
Further denote (t.b)
Pir.
Pj.n. P(pj
(t.h)
(1.b)
0)
pltab)
Phra
(t.b)
(P 0). Then
(t.b)
j.r
w (1.b) Jir = W J.r + jijir ||11|12
(0.0)
(t.b)
(t.h) Pain
E: + p
((. 11:11 (7)
P
Note that (7 is a variant of 66) by decomposing the coefficient Pj.n
(t.b) into Parr and e e we
(t.h)
can streamline our proof process. The normalization terms p-.. || and ||E:112 ensure that (1.6)
for ..i). Through signal-noise decomposition, we (t.b) and the learning progress of noise using Pj.r This decomposition turns the analysis of SGD updates into the analysis of signal noise coefficients. KKou et al. (2023) extend this technique to the ReLU activation function well as in the presence of label flipping noise. However, mini-batch SGD updates amplify the complications introduced by label flipping noise, making more difficull 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 T5.r update rule is as follows:
Tjrr Rs (wfth)) .1) and Pir (P () (1.)
characterize the learning progress of signal la using 7j.r
(t,b) whose
(1,++1) (1.b)
17
Σ
(1.b) d'(w ;0 12))
7jr = -7i
Bm
VEI,ANS+
Σ
(((i)) d' d'ww
(1.b) -0 mi) |1e11/2.
(8)
iEI
Here, Its represents the indices of samples in batch of epoch t, S. denotes the set of clean samples where y, Di, and S_ represents the set of noisy samples where y1 = - The updates of (1.b)
(f.b)
Tjrr 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 Eietinn e((c))) of (WW UV 14) primarily drives model learning. Such unfavorable scenarios ('ww)) (t.b)
can be attributed to two possible factors. Firstly, the gradient of the loss e (((t.)) might be significantly larger for noisy samples compared to clean samples. Secondly, during certain epochs, the majority of samples may be noisy, meaning that I, S__ significantly outnumbers I,, nS++
To deal with the first factor, we have to control the ratio of the loss gradient with regard to different samples, depicted in (9. Given that noisy samples may overwhelm a single batch, impose an additional requirement: the ratio of the loss gradient must be controllable across different batches within a single epoch, i.e.,
(1bb))) e (t.bba) S SC2.
As ((1))(((2)) As exp(z2 - 21), we can upper bound g((.be) f(W((.b2) xx), which can be further upper bounded by Σ, Paa
(t.b)) jel (g(ttbb)) by yi f(W(t.b)) x;) 1/4 (t.bg)
Erpli purkk with small emror. (t.b)
Therefore, the proof of (9) be reduced to proving uniform bound on the difference among P (t,b.) i.c., DL.r.
(1,b2) Puink < K, Vi
However, achieving this uniform upper bound turns out to be challenging, since the updates of (1.M). are not evenly distributed across different batches within an epoch. Each mini-batch update utilizes
only portion of the samples, meaning that some P can increase or decrease much more than

---

the others. Therefore, the uniformly bounded difference can only be achieved after the entire epoch processed. Consequently, we have to first bound the difference among P(..) after each entire epoch, and then control the maximal difference within one epoch. The full batch (epoch) update rule established as follows:
n(P -1P--13. 1)2
((+1.0) Parr Pu..nkk
(1.0) Pyira
Pan.n.
115.12
Bm
B ) 1 10
((1.) 11Ex11)
(10)
Here, (/() denotes the batch to which sample belongs in epoch t, and S 3(0,5)) parameters that learn &; at epoch & defined as
represents the
3/(.b) fr (w (6.0) Besr
E.) 0}.
(11)
Therefore, the update of > [P...) - I is indeed characterized by the activation pattern of parameters, which serves as the key technique for analyzing the full batch update Σ P
(t.b) p However, analyzing the pattern of s(t.b) directly is challenging since (w g.."" E.) fluctuates
in mini-batches without sample i. Therefore, we introduce the set series S/t.b) as the activation pattern with certain threshold follows:
s((:0) {r: (w(t.b) Wiir Ei) > doop Vd/v///
(12)
The following lemma suggests that the set of activated parameters s(t.0) is non-decreasing sequence with regards to i, and the set of plain activated parameters 5(C.0) always include s(t.0) Consequently,
s(U,,) is always included in (((..) guarantecing that &, can always be learned by some parameter. And this further makes the proof for SGD, we consider the learning period 0 I t T", where T = 1 poly(c d, n, m) is the maximum number of admissible iterations.
the difference among Py.rt is bounded, well (((t.11) let /(((.bb))) < C2. In
Lemma 3.4 (Informal Statement of Lemma C.8). For all t E [0, and b E [H. we have
S ((00)00) C (1.0)
5 (((b))
(13)
As we have mentioned above, if noisy samples outnumber clean samples, (t.b) (t.b) 7j. may also decrease.
To deal with such a scenario, we establish two-stage analysis of "irr progress. In the first stage, when - is lower bound by positive constant, we prove that there are enough batches containing sufficient clear samples. This is characterized by the following high- high- probability event.
Lemma 3.5 (Infomal Statement of Lemma B.6). With high probabilit,, for all T E (0(1), there exist least ej T epochs among [O, T'. such that at least e2 И batches in each of these epochs satisfying the following condition:
|S+ Cunn S. II,6 E 10.25B 0.75B).
(14)
After the first stage of
2d-1 epochs, have Tj.r "(P-1)" (t.b)
(T.0) Ω
|
The scale of (T,0) ?/r guarantees that (1.) μu) remains resistant to intra-epoch fluctuations. Con-
sequently, this implies the sign of (5.b) 14) will persist unchanged throughout the entire epoch.
Without loss of generality, we suppose that (w jor (x) > O, then the update of //.r can be written as follows:
(1.b)
(t,b)
(t+1,0) = =1 Jrr Bm 255
(t.0)
77
min iEI....b
1 (t.b)
||S nsi max (1.b) i62,,..b
| 1Is_|ssnss
ns_.]
|1 (15)
As we have proved the balancing of logits (((,)) across batches, the progress analysis of ((+1.0) is established to characterize the signal learning of SGD.
7ir
7

---

Result for SAM
In this section, we present the positive results SAM in the following theorem.
Theorem 4.1. For any O, under Condition 3.1 with do
choose PopVa) With probability at least - neural networks first trained with SAM with o t "mlppll))) iterations, then trained with SGD with ((1=-t--mmdd 2-22pp" -
θ
myBB
iterations can find w(() such that,
1. The training loss satisfies (.((W())) Se.
2. The test error -
(W()) sp
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, opposed the threshold of || 2/2 2 (2((d/))PP)) for CNNs trained by SGD. The discrepancy in the thresholds can be observed in Figure [ This difference is because SAM introduces perturbation during the model parameter update process, which effectively prevents the early- memorization of noise by deactivating the corresponding neurons.
4.1 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, have the following update rule of decomposition coefficients
(1.bb) (5 P 3.3
Lemma 4.2. The coefficients (t.b) (t.b) ((t.b) defined in Lemma 3.3 satisfy the following iterative .p5tr
Vir Pj Lirz
equations for all E [m], j {=1} and [n]:
Yir .Pj.r. Lyrr. = 0,
(0,0) (0.0)
(1,++1 (1,b) Fjr -ij Bm
(lbb) of ((w (1.0)
+ jjr .0i 14)
iEI,,nSs
(1.0) ((t.b) o' d'(( 10r
+ =(t1))
(i))
11m112
VEI,
(1,b1)) (t.b)
()P 1)2 ((..b) o' ,''(w + +GJ,r E.)) ||8.11 11ys = j) 1(i Ic,o,,
(1.0)
(t.b)
Pard
Pj.ri
Bm
(t.b)) Lir +
1/P 1
p(t,b)
d'((w + 2.r
(l.b)
((.b)
(.) ||E:112
[1 = -j) l E I,,)),
Bm
where Ie,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 w (r.b) <(t.5)
j.r fjr whereas in SGD, it is determined by
(t.b) This perturbation to the weight update process at each iteration gives the unperturbed weight Wir
SAM an intriguing denoising property. Specifically, if a neuron is activated by noise the SGD update, it will subsequently become deactivated after the perturbation, 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 (w (t,b)
(t.b)
Ex) 2 0, k E Iz,b and j then (w
(. (t.b)
Ex) < O.
jr
j.r
jrr
By leveraging this intriguing property, we can derive a constant upper bound for the noise coefficients (t.b) Pana by considering the following cases:
(1.b) 1. If E is not in the current batch, then P will not be updated in the current iteration.
8

---

1.0
10
Dimension
Dimension
0.6
0.6
11000
0.4
0.4
1000
0.0 1000
0.0
Strength of the Signal
Strength of the Signal
(a) Original SGD Test Error Heatmap
(b) SAM Test Error Heatmap
Figure 2: (a) is a heatmap illustrating test error on synthetic data for various dimensions d and signal strengths la when trained using Vanilla Gradient Descent. High test errors represented in blue, while low test errors are shown in yellow. (b) displays heatmap of test errors on the synthetic data under the same conditions as (a), but trained using SAM instead with 0.03. The y-axis represents a normal scale with a range of 1000 ~ 21000.
2. If E, is in the current batch, we discuss two cases:
(a) If (w (1,8)
Jar {i) 2 O, then by Lemma 4.3 one can know that o' d'w
(t.b)
jir + 2(t.)) .E:) = 0 and
(t.b)
thus Pjan will not be updated in the current iteration.
(b) If (w (c.0) 3.
{i) S 0, then given that (t.b) ..i = (1.b) and ((l.b11)
(t.b)
Jrn + 1)
S
Bm
we can assert that, provided is sufficiently small, the term p can be upper bounded by small constant.
In contrast to the analysis of SGD, which provides upper bound for Pj.r.i of order Oloog d), the noise memorization prevention property described Lemma 4.3 alloows us to obtain an upper ((.bb) bound for Pj.r.i of order O(1) throughout [0, Ti]. This indicates that SAM memorizes less noise compared to SGD, On the other hand, the signal coefficient Tj.r. also increases to Q(1) for SAM, following the same argument 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.
(t.b)
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 invariant, without loss of generality, we set f4 = |A|0| 0,.0...., O We then generate the noise vector & from the Gaussian distribution (O o0 with fixed standard deviation op 1. We train a two-layer CNN model defined in Section 2 with the ReLU activation function. The number of filters is parameters and train the CNN with full-batch gradient descent with a learning rate of 0.01 for 100 iterations. We consider different dimensions ranging from 1000 20000, and different signal strengths |/1// ranging from to 10. Based on our results, for any dimension d and signal strength setting we consider, our training setup can guarantee training loss smaller than 0.05. After training, we estimate the test eтror for each case using 1000 test data points. We report the test error heat map with average results over 10 runs in Figure
20 without label-ftipping noise. Since the learning problem is rotation-
m 10. We the default initialization method in PyTorch to initialize the CNN
Related Work
Sharpness Aware Minimization. Foret et al. (2020), and Zheng al. (2021) concurrently introduced
methods to enhance generalization by minimizing the loss in the worst direction, perturbed from the current parameter. Kwon et al.) (202T) introduced ASAM, variant SAM, designed to address parameter re-scaling. Subsequently. Liu et al. (2022b) presented LookSAM, a more computationally

---

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 viewed as generalization of the SAM algorithm. Meng et al. (023) 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. [i et al. (202 b); 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. [ao et al. (2022) provided an analysis for learning two-layer convolutional neural networks (CNNs) with polynomial ReLU activation function (ReLU%, g > 2). Kou et al. (2023) further investigates the phenomenon of benign overfitting in learning two-layer ReLU CNNs. Kou et al. (2023) most related to our paper. However, our work studied SGD rather than GD, which requires advanced techniques 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, 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, YK, and QG are supported in part by the National Science Foundation CAREER Award 1906169 and IIS 100990,,,,, and the Sloan Research Fellowship. XC and CJH are supported 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, K., BUBECK, S., CHEWI, S., LEE, Y. T., SUAREZ, F. and ZHANG, Y. (2022). Learning threshold neurons via the" edge of stability". arXiv preprint 22X.v:2212.07469.
ALLEN-ZHU, Z. and LI, Y. (2020). Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint arXiv:2012.09816.
AL..EN-ZHU, 7. 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, M. and FLAMMARION, N. (2022). Towards understanding sharpness-aware minimization. In International Conference on Machine Learning. PMLR.
ANDRIUSHCHENKO M., VARRE, A. V., HILLAUD-VIVIES L. and FLAMMARION, N. (2023). Sgd with large step sizes learns sparse features. In International Conference on Machine Learning. PMLR.
10

---

BAHRI, D., MOBAHI, H. and TAY, Y. (2021). Sharpness-awe aware minimization improves language model generalization. arXiv preprint arXi::2100.0859.
BARTLETT, L., LONG, P. M. and BOUSQUET, (2022). The dynamics of sharpness- minimization: Bouncing across ravines and drifting towards wide minima. arXiv preprint arXiv:2210.01513.
BARTLETT, P. L., LONG, M., LUGOSI, G. and TSIGLER, A. (2020). Benign overfitting in linear regression. Proceedings of the National Academy of Sciences.
BEHDIN, K. and MAZUMDER, R. (2023). Sharpness-ave aware minimization: An implicit regularization perspective. arXiv preprint arXiv2230..1183.
BEHDIN, K., SONG, Q., GUPTA, A., DURFEE, D., ACHARYA, A., KEERTHI, S. and MAZUMDER, R. (2022). Improved deep neural network generalization using m-sharpness-awwe minimization. arXiv preprint 2:Xi:2222..04445
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, 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, N. S. and LONG, P. M. (2022). Deep linear networks can benignly overfit when shallow ones do. arXiv preprint 2091:52009.0931..
CHEN, X., HSIEH, C.-J. and GONG, (2021). When vision transformers outperform resnets without pre-training or strong data augmentations. arXiv preprint arXiv:2106.01548
DEVROYE, L., MEHRABIAN, A. and REDDAD, T. (2018). The total variation distance between high-dimensional gaussians. arXiv preprint arXiv:/810.08693.
FORET, P., KLEINER, A., MOBAHI, H. and NEYSHABUR, B. (2020). Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412.
FREI, S., CHATTERJI, N. S. and BARTLETT, P. (2022). Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. Conference on Learning Theory. PMLR.
JELASSI, S. and LI, Y. (2022). Towards understanding how momentum improves generalization in deep learning. In International Conference on Machine Learning. PMLR.
KOU, Y., CHEN, Z., CHEN, Y. and GU, Q. (2023). Benign overfitting for two-layer relu networks. arXiv preprint 453.42330..04145.
K KWON, J., KIM, J., PARK, H. and CHOI, I. K. (2021). Asam: Adaptive sharpness-aware minimiza- tion for scale-invariant learning of deep neural networks. In International Conference on Machine Learning. PMLR.
LI, Y., WEL, C. and MA, T. (2019). Towards explaining the regularization effect of initial large learning rate in training neural networks. In Advances in Newral Information Processing Systems.
LI, Z., WANG, T. and ARORA, S. (2021a). What happens after sgd reaches zero loss?-a mathematical framework. arXiv preprint arXiv:21000061144
LI, Z., ZHOU, Z. and GRETTON, (2021b). Towards an understanding benign overfitting in neural networks. arXiv preprint arXiv:2/06.03221.
LIU, Y., MAI, S., CHEN, X., HSIEH, C.-J. and You, Y. (2022a). Towards efficient and scalable sharpness-awan aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).
LIU, Y., MAI, S., CHEN, X., HSIEH, C.-J. and YOU, Y. (2022b). Towards efficient and scalable sharpness-awar aware minimization. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition.

---

MENG, X., CAO, Y. and ZOU, D. (2023). Per-example gradient regularization improves learning signals from noisy data. arXiv preprint arXiv:2303.17940
MONTANARI, A. and ZIIONG, Y. (2022). The interpolation phase transition neural networks:
Memorization and generalization under lazy training. The Anmals of Statistics 50 2816-2844
OLSHAUSEN, B. A. and FIELD, D. J. (1997). Sparse coding with overcomplete basis set: A strategy employed by v1? Vision research 37 3311-3325.
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.0157.
VERSHYNIN, (2018). Digh--Dimessioal Probability: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics, Cambridge University Press.
WEN, K., MA, T. and LI, Z. (2022). How does sharpness-aware minimization minimize sharpness? arXiv preprint arXiv:2211.05729.
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, 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, Y.. ZHANG, 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, 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.0800 08065.
12

---

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 76.4%.
10%
30 %
50%
%
90% %
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.3
76.2
Table Top-1 accuracy (%) of wide ResNet on the CIFAR datasets when adding different levels of Gaussian noise.
Model
Noise
Dataset
Optimizer
Accuracy
WRN-16-8 WRN-16-8
CIFAR-10
SGD
96.69
CIFAR-10
SAM
97.19
WRN-16-8
N/O0 0.1)
CIFAR-10
SGD
95.87
WRN-16-8
((0,0.1)
CIFAR-10
SAM
96.57
WRN-16-8
N/O,
CIFAR-10
SGD
92.40
WRN-16-8
0.3,
CIFAR-10
SAM
93.37
WRN-16-8
N(0,
CIFAR-10
SGD
79.50
WRN-16-8
((01
CIFAR-10
SAM
80.37
WRN-16-8
CIFAR-100
SGD
81.93
WRN-16-8
CIFAR-100
SAM
83.68
A Additional Experiments
A.1 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 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 10k steps linear warmup, then cosine decay. As shown in Table the earlier SAM is introduced, the more pronounced its effectiveness becomes.
SAM with additive noises. Here, conduct experiments the CIFAR dataset with WRN-16-8. We add Gaussian random noises to the image data with variance (0.1, 0.3. }. We choose the hatch size 128 and train the model over 200 epochs using learning rate of 0.1, momentum of 0.9, and weight decay of 5e - 4. The SAM hyperparameter is chosen as r = 2.0. As we can see from Table 2 SAM can consistently prevent noise learning and get better performance, compared the SGD, which varies from different additive noise levels.
A.2 Discussion on the Stochastic Gradient Descent
Here, we include additional empirical results show how the batch size and step size influence the transition phase of the benign/harmfüs regimes.
Our synthetic experiment is only performed with gradient descent instead of SGD in Figure We add an experiment SGD with a mini-batch size of on the synthetic data. we compare Figure and Figure the comparison results should remain the same, and both support main Theorems 3.2 and
In Figure we also conducted an extended study the learning rate to check whether tuning the learning rate can get better generalization performance and achieve SAM's performance. Specifically, experimented with learning rates of 0.001, 0.01, 0.1, and under the same conditions described
13

---

1.0
1.0
0.6
08
Dimension
0.5
Dimension
0.
11009
0..4
04
-0.2
02
1909
0.0 1000
0.0
Strength the Signal
Strength the Signal
(a) SGD Test Error Heatmap with minibatch 10
(b) SAM Test Error Heatmap with minibatch 10
Figure 3: (a) is a heatmap illustrating test error synthetic data for various dimensions d and signal strengths jJ when trained using stochastic gradient descent. High test errors are represented blue, while low test errors shown in yellow. (b) displays heatmap of test enrors the synthetic data under the same conditions as in (a), but trained using SAM instead. The y-axis represents a normal scale with range of 1000-21000.
in Section [ 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 perfor- mance. 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 eL et al. 2019 2021a Ahn et al. 2022; AAnriishhhhenoo et al., 2023). Importantly, even with this expansion, the benign overfitting region remains smaller than what empirically observed with SAM. Our conclusion is that while SGD, with more considerable learning rate, exhibits improved generalization, it still falls short of matching SAM's performance.
B Preliminary Lemmas
Lemma B.1 (Lemma B.4 in Kou et al. (2023)) Suppose that and 2[log(6n/ /6). Then with probability at least 1 d,
d/ < 1651 < /
(E:, Ex)1 s 20 /dlog(6n2/0),
(E m)l ||u||aap /2log(6n/00
for all i, [[n].
Lemma B.2 (Lemma B.5 in Kou et al.
2023)
Suppose that d = 2(log(mn/8)), m 2(log(116))
Then with probability at least
odd/2
(0.0)
1| 3
j.r
((w j.r 1)) 2 12m/6)
(0,0)
1|1,,
(( (w,0)) &i)l < 2 eoy/og(12mm)) 0oop vd
for all E [m], { and i E [n]. Moreover,
aol|m|12/22 < max j (w (0.0) rE(m]
j.r 14) /
c0l|u112.
0oo vd/4 S max j (w 0.00)
EE) s 2/log 12mn/5) Vd rE[m] jr
for all j E {+1} and [n].
Lemma B.3. Let s((.b) denote fr (tb) E.) ooo, vd/v2///2)3. Suppose that 5 and 2 50 (0)(()))))). Then with probability at least 1 d,
|S(0.0)
0.84 1mm, Vi E [n].
14

---

10
1.9
08
0.8
Dimension
Dimension
0.6
11000
11900
0.4
0.4
5.2
0.0
Strength of the Signal
Strength of the Signal
(a) Full-batch GD Test Error Heatmap with learn- ing rate y/ = 0.001
(b) Full-batch GD Test Error Heatmap with learn ing rale 1) = 0.01
1.0
1.0
0.
0.8
Dimension
Dimension
L1000
11000
0.4
0.4
-.2
0.2
1000
Strength of the Signal
Strength of the Signal
(c) Full-batch GD Test Error Heatmap with learn- ing rate 7)
(d) Full-batch GD Test Error Heatmap with leam- ing rate !)
Figure 4: (a) is a heatmap illustrating test error on synthetic data for various dimensions d and signal strengths a when trained using learning rate 0.001. High test errors are represented in blue, while low test errors are shown in yellow. (b) displays 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, The y-axis represents normal scale with a range of 1000-21000.
Proof of Lemma
(0.0) B.3 Since (w yeir
Es) ~ (O, of & 112), we have
P(( yer
(0.0)
E.) > oo Vd d 2) 2 P( (0.0) u1.
&
a Foll):112) = =000=-(11 ((-1).
where ((-) is CDF of the standard normal distribution. Note that |s!0.01| = Em= 1[( (W < >
ooo,, Vo vd/v2 probability least 1 8/n, have
and P(w ye.r"
(0.0)
E.) oOo, vd/v2/ 2 (--1), then by Hoeffding's inequality, with
IS (0.0)
log (2n/8)
(--1)
m
2m
Therefore, as long as 0.2 0m2
2 1og(2n/5)
by applying union bound, with probability at
least 1 - 5, we have
|s 2 0.84( Vi [n].
Lemma B.4. Let s((.))) denote (i E [n] : 3r j. (w {i) > Gooa odo,Va/V2), Suppose that 5 and n 2 32 3ogg(44m/)) Then with probability at least
(0 m Vj E {+1), [m].
15

---

Proof of Lemma B.4 Since (w (0.0)
E.) ~ (O (1 we have
j.r
P((w)r 1 E ,oo, vd/v2) P((w), (0.0)
(w E.) gnllE:1l2) +((-))
where 4D-) CDF of the standard normal distribution.
Note that (51 = Ei=1 1i = 5 [[ Ei) > ouoVdd/V2 and P(y. = j. (w Ei) > poo,vd/v2) (211 then by Hoeffding's inequality, with probability at least 8/2m, we have
log (4g(4m/6)
((-1)/2 +
2n
Therefore, as long as ((-1)/4 at least 1 1- 6,
log( (1mmmm 15) 2n
by applying union bound, we have with probability
n ( VJ (+1), e [m]]
Lemma B.5 (Lemma B.3 in Kou et al. 2023) ). For |S+ SySyl and |S_ n Syl where y {+1},
holds with probability at least 1 - o(d > 0) that
p)n
n log (c). 1S- nsy pn S V/2 log (9) Vy Eyyeet)) {+1}.
[S+ S
2
2
Lemma B.6. It holds with probability at least 6, for all T E [log(2T" (0) (6) T and E {+1),
cg.
there exist at least cg T epochs among [0, T.. such that least C4 H batches in these epochs, satisfy
B 3B
1S+n n II,al E
(16)
4 4
Proof. Let
n
E1,t == {In epoch 4, there are at least c2
batches such that 16 holds for y = 1},
B
E1,46 = {I [In epoch natch b. 16 holds for y ==
First let n big enough, then have S+ Sy [_(1=) 5(1 [[]-p]n] We consider the first ei H 8 batches. At the time we are starting to sample fi-th batch in the first ej H batches, suppose there are n1 samples that belong to S4 n Sy, and there are n2 samples that don't belong S+ Sy. Then
n1 2 3(1-p)m
5(1-p)m and n2 2 3(1-p)
5(1 5(1-p)n
Gin 2
Gin 2
16
8
16
3B/4 Ei=B/4 C CEC, C.B
P(E
ai
CB
2i=1B/4 C pin
3B/4 Cl
Cl
CB- pin
CB
B BCB/
B/4
11/4(911111 ((1-p))
B BBBB
B
32
n
B B/1
9(1 - p)
B
202-
32
Then, the probability that there are less than C10HH batches in first e₁ H batches such that 16 holds is:
Σ
Σ
P(1(E)10)= = 4o (E),.))) = 41 (E1,0) = lol... ...
1=0
Eh=
16

---

= le, H-2 P[1(E1,t,cHH)) = H- (( = ... (( -
6 Σ C qn -202)
0-0
S (202)
Choose Ho such that ej Ho (20g) ) (1 202) jey = 1- 2c3, then as long as H 2 Ho. with probability C3, there are least C102 batches first H batches such that 16 holds. Then P[E1,1] 2 20g. Therefore,
2t
P Σ 1(E,) 2Tc3 S -t) exp(-
Let T 2
log(27" /8) 2cj
Then, with probability at least 8/(2T*),
I 2 csf'
Let C4 = C102. Thus, there are at least csT epochs, such that they have at least o,H batches satisfying EquatiorB.6 This also holds for 1. Taking union bound to get the result.
C Proofs for SGD
In this section, we prove the results for SGD. We first define some notations. Define H = -n/ as the number of batches within an epoch. For any & b2 and bj, by [H] we write (t1. b1) S (t.b) < (t2, b2) denote all iterations from t1-th epoch by-th batch (included) fz-th epoch's by-th batch (included). And the meanings change accordingly if we replace < with <.
C.1 Signal-noise Decomposition Coefficient Analysis
This part is dedicated to analyzing the update rule of Signal-noise Decomposition Coefficients. is worth noting that
M
1 m
F/W,XX x)
ΣΣ
a(w J.r. xp))
Σ o(w B1)) + (P- (P)))o(((,,r E)) m
Let Ir, denote the set of indices of randomly chosen samples at epoch batch b, and [Z,z| = B, then the update rule is:
for b H
W (1b+1)
(1.b)
Vw5,, W... LIA (W(+.b)
V
I1(P 1)
-w(1.b)
((e.))
o' o'((w ((.b) ..i) jy
J.r
Bm
IEI..b
1 Bm
(1.b)) o'(( ( jv.D...
(t.b)
TET,
(++1,0)
(t.H)) j.r
and
(17)
C.1.1 Iterative Expression for Decomposition Coefficient Analysis
(1.b) ((t.)) p(t.)) defined in Lemma 3.3 satisfy the following iterative Lemma C.1. The coefficients Vi Phnn fira
equations:
(0,0) (0.0) (0,0)
Ojr : Pi.ri 1 j.r.
0,
(18)
(1,+1) (1,b)
17 Bm
Σ
(LD ((b)) d(
(1.b)
p))
i
17

---

Σ (:a d'((''
((
994 mb)) 11
(19)
NET,,ONS_
Pjnú = P/101
(lbb1))
n(P 1PP-10 1 (((.b) o' d'((
(1.0) .E() -1E::13 1(v: 5) I( I,bb,
(20)
Bm
n(P- 1) ((t.b) o' d'((w)) (1.0)
fjnd ((ottt)) = (1.b)
.E)) -1E::12 1 = -i) E I(b),
(21)
Bm
for all E [m], E {=1} and E [n]]
Proof. First, iterate the gradient descent update rule epochs plus b batches and get
(1.b) orr
(0,0)
(P
1)
E Σ
(b) (V o' d' j.r E.)) (P 1);,
w jir
Bm
(.
π
Σ
E
((w! ()) (.il)) y.hiju.
Bm
(M.D)C(LL))) iEZ,,
According to the definition of 7jr and (t)
Pj.r.
3.
(8.6)
(0.0)
+ 735 75
(1.b)
||
+
P
(t.b) Pjr.
||E:112 EE.
Since &, and μ are linearly independent with probability 1, we have the unique representation as follows:
(t,b)
n(p- - 11)
Σ
(re o'((w) Jir E.)) -1E: (i E Ia,))ui,
(1'8) ((W(('.)
Bm
(W(D))((t)))
(t.b)
17.
Σ
Ml 0' d(ww (w 33 14)
Bm
(M,)))((t)))
pl o' d' yi (o)) 1111..
iEI
Since we define 0 = Pjrn 1(p)
(1,00 2 0), (t) = (1.)) (1.0)
S.r.s Piin 1( S 0), we obtain
(1.b)
111
Σ (( d' d'(WW)r jrr (i) | 1(i E 1e j),
(1
Pj.r.
Bm
(P,))<(()))
(t.b)
(PP-1 1)²
() d ((w
8' (.) 1 (i Te, 1(y = -j).
Bm
((t.))((t.))
And the iterative update equations
19 20 and 21) follow directly.
C.1.2 Scale of Decomposition Coefficients
We first define T = =1 "poly( d, n, m) and
a = 4 log(TT))
(22)
B = 2 max {l
iJ.t j.r
(0.0) .1).. (P (P-1)|(W 1)(W (0.0)
(23)
||1/|2
SNR
(24)
(P- L)
By Lemma B.2 and Condition 3.1 can be bounded as
8 = 2 max
f...r jrr 14), (P - 1) (ij E
(((
(0.0) (w
(0.0)
18

---

S 2 max{ 21lgo020g11m oo F011 log (12mmmmm))))))) -0p--ooPPpooo0 - 1)0pVd}
= O (lo((mn/0) ao 1) dd,,
Then, by Condition 3.1 we have the following inequality:
32 log(6n/0) na,5p 5V log(6n2/0)
1
max p, /
1 12
(25)
d
d
We first prove the following bounds for signal-noise decomposition coefficients.
Proposition C.2. Under Assumption 3.7 for (0,0) s (1,b) s (T*, we have that
3.r Pj...i Li
(0,0)
(0.0) (0.00 0
(26)
0 < Pj.ra S az,
(1.b)
(27)
1og(6n2/0) (5)
0 ((t 3 100
na a.,
(28)
and there exists a positive constant C" such that
(1,b) S Tjr S C'Fa.
(29)
12
for all r E jm] {=1} and i E [n]), where 7 n SNR?.
We will prove Proposition 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.7 suppose
and hold after b-th batch of I-th epoch.
Then, for all E [m]] E {+1} and E [n].
(w!lho)
wol. 3.r 14) - -33
(0.0)
(t,b)
I SNR
32 (0g(6//5)
na.
(30)
d
(0.0)
-E)
(ft.h)
5
log 6/8/8)
na, # yss
(31)
p 1
P
d
w jrr
(0.0)
Ei)
1
(t.b)
<
5
log((n2 (8)
na, y/a
(32)
P
P
d
Proof of Lemma by definitions,
C33
First, for any time (t, 2 (0, 0), we have from the following decomposition
jr
(t.6)
w jjr 14) =
(0.0)
(t,b)
P
1
(1,b)
11222 2.
. (Ev,HL)
+ p
(1.b) Lina 1r: 2 -(v,
U=1
According to Lemma B.I we have
1
Parn |1E 115: (Ev, 1))
+ P-
p (4,b) June
11 (Ev. p4)
P 1
7=1
<
l0 (2.0)
0 II|E 11
1(5 14))+ +
Sur |1|s:r 11 1KEv. 14)1
(o,b)
111::-11122
P
P 1
1=1
==1
2| 11/11/22 log 0
n
<
+
(P- 1)
19

---

Blo((6n// (6n/o)
72
V
Σ
(c:b)
+
===1
32 log (6n/8) na;
SNRV
d
where the first inequality is by triangle inequality, the second inequality is by Lemma B.1 the equality
is by SNR = |1ll2/((P - 1) dd, and the last inequality is by 27 28 It follows that
(e.b) (w() jr
(0.0) j.r 1)) -j -7500)) S <SNRV
32 log 66 6n/0)
na.
Then, for j + yr and any & O, we have
(w (1.b)
w jrr Et)
(0.0)
1
= j /3.r ||
{M,. Ei) + P-1 1
Panne Paneeller |15r112? (Ev, Ei)
1
(C.B
|1 2.
(E)
j Jir || 2 (pe. Ei)
(t.b)
+ p-1 1
1
|15:1122 (Ev,
1 (t.b) P- J.r.
+j
Ti.r
(t.b)
|el12²² (M,, Ei) + P
(1bb) Sin
|1 11a
(E2, E1),
where the second equality is due to (t.b)) j.rf
0 for j + yi. Next, we have
(1.b) 3 7j. || (1, E.)
1
|1 l1
(E),
p
pzz
$ 17
(0 1|1eo112 2. 114 Eiحl +
1e -KEr,
p
log(6n2/0)
(1,00 s 17 1| /2 log(6n//))
+
p -
d
SNR 1
2 llg(///
(2(00)
4
log (6n (6n2/8) /8)
(ttbb) J,r,""
p P-17
P-1 1
d
SNR
8C log(Bn//))
4
log (62//)
P 1
d
na +
p
d
na
log(6n2 15) P
5
d na.
where the first inequality is by triangle inequality; the second inequality is by Lemma B.1] the equality is by SNR = || d: the third inequality is by (28) and (9); the forth incquality is by SNR < 1/V8 Therefore, for j * ye. we have
(1,b) (
J.r
(0.0) Ex)
1
(1,) Sin
S
5
log (6n (6n2/8) na.
P
p
d
Similarly, we have for yi j that
(w (e.b) w j.r E
(0.0)
1 n
=j - (1.b6) ||111122 (1,, Ei)
(t,0)
|1::.11 (Ev, Ei)
p
1
#=1
20

---

1
|| - Ei)
P
= =j. ||11|12
(1.b)
(1e, Ei) + p 1
(1.b) Pung
|5
(Ev, Et)
8
p P-1 (6-0) + T,.
1
(c.b)
1111 2 (M, Ei} +
Σ
Sirasr
||E: 112 Ei),
p
and
(2.00
1| (M,, E.) +
Σ
Plon ||5: (Ev. Ei)
8
SNR 1
2 log (ogg(6n/))
(t.b))
4
log(6n2/0) 15)
$ P P-1
+
d
SNR
8C2 log(6n/00)
log (6n /6) d
P-1
d
na +
P
na
5
log(6n²/00 na,
P-1
d
where the second incquality is by Therefore, for j = yi, we have
(7 and 29 and the third incquality is by SNR & 1/V80"2
((wfer) (1.b)
(0,0) Ei)
1 (1,b)
5
(g((6n2/))
J.r
P-
j.na
S
P
d no.
Lemma C.4. Under Condition 3.7 suppose 27 p.
(28) and (29
hold after b-th batch of t-th epoch.
Then, for all j F yo E [+1] and i E [m), (W .x.)
S 0.5.
Proof of Lemma C.4
According to Lemma
we have
1 m
(W (1,b) j
x,)
((
jo U.pp)) + (P 1) (1.)) j.r .E)))]
(A.b)
m
max
(1,)) (w Jir ..:H)) (P- (w (1.b)
S max (w vit)), (P 1w
(0) E,), SNR
32 log (6n/d)
no, yij (t,b)
log (6662/2/)) na + (((.b)
5V
d
j.ri
32 log(lo//)) 6n/0)
(1.b) log (6n2/))
< 6 max
///2
na, -7gir 5/
na
d
d
0.5,
where the second inequality is by 30 and p(C.)) jaz < 0; the third inequality is by (5) and
3 and
32
;
the third inequality is due to the definition of B (t.b)
?j.r
Lemma C.5. Under Condition 3.1 stiposs Then, it holds that
27 and
29
hold at b-th batch of t-th epoch.
(P - 1) (1,b]
<i) 2 -0.25,
Ne.r
(P- 1) (f.b)
(1.0)
Ei) S (P 1
1)
((w
(Lb) Mi.r
E.)) S (P
(1.) We.t Ed) + 0.25,
We.r
for any i E [n].
21

---

Proof of Lemma C.5
According to 32 in Lemma C.3
we have
log(6n2 /6)
(P- 1)
(t,b) 2arr
Et)
(P 1)
(U,0) y.r"
Ei)
Poi.r V
d
llo((nn /8)
2 3 5m
2 0.25,
where the second inequality is due to Pu.
O, the third inequality is due to B < 1/8 and
5n /0gg < 1/8 by Condition 3.1
For the second equation, the first inequality holds naturally since S o(2). For the inequality, if (Wyer: Ei) S O, we have
(t.b) (P- (--))(((W(.)) Harr Ei) 5 (P- 1)( (1.) E.) + 0.25.
And if (Wiirrs Et) > O, we have
(1.b)
(P ((w pirt Ei) - (P
(1.b)
Wiet Es) < (P- 1)(w()))) (1,))
Miit Ex) + 0.25.
Lemma C.6 (Lemma C.6 in Kou al. (2023) ) Let g(z) (())) 11(11 + exp(z)), then for all 22 2 2 -1 where 2 O we have that
exp(c)
9(21)
exp(c).
g(22)
Lemma C.7. For any iteration E [0) and b, b,, b2 E [H) we have the following statements
hold:
1.
E (1.0)
Puint
(t.0) -r
((.,b)) [(hbba) Puint Pannk
< 0.1x.
2. (w!fit) (1.b)
Ei) 2 (w wur Eii 0oopVa/V2.
(f.U)
3. Let (((.b) fr E [m] (w/lt)) Ei) > 0), then we have
(1.b)
s(r.0) C (((b))
4. Let (((.b) jr = {i E [m] yı = j. (.(w(( (L,b)
Es) > 0}, then we have
jor
s(t,,)
((t.b) 3,r
C
Proof. For the first statement,
Pul.rs Pun.rc
1(P 1)2
( (
11E, (s((-11b2) | 11
||Ex 113
max
Bm
(P 2)a
303d
B
2
0.1x,
(t.b)
where the first incquality follows from the iterative update rule Pann 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
(t.b) Ma.
E = (w ys.r
(t.b-1)
Ex)
17.
a (l,b-1) o' ((w.t.b- (t.b-1) 1), (.))))) (yv us, E.
Bm
FEI,a-1
22

---

T(P-1)
Σ
((b-11) d' d'(wltib"), Ei).
Bm
PETLa 1/1
Therefore,
(w(tt)) (1.6) Birr
(w.).)) (t.0) Nir
(t.0) E.)
17
|Ill|22 2log(6n/0)
n(P- -P-)) 1) 20 d
Bm
Bm
{,) do0p v 4/v2.
where the first inequality is due to Lemma
B. and the second inequality is due to Condition 3.1
For the third statement. Let r* E s((00) then
(w (t.b) Best
(r.0) E.) - Joo, VoopVa//22>>
where the first inequality is due to the second statement, and the second inequality is due to the definition of 50,00 Therefore, r* E S((..) and s(t.b) C ((C.b)) 'he fourth statement can be obtained similarly.
Lemma C.8. Under Assumption 3.1 suppose (27
28 and 22
hold for any iteration (L', 8) s
(1, 0). Then, the following conditions also hold for Vi' I and 7 8/ b/ E [H]:
E ((t-00) - ((,00) S K. for all i, k [n].
2. -f(W((.))) xx) Uh f( f(w (M, xx) S for all i, k E [n].
3. et ((EM))
M))
/1th (116.66) 1b) C2 exp(C1) for all i, [m].
s s(0.0) c s((10) where 5(1'10) = {r E [m]] (1'-0) , ,E)) > voovVa/V2)) and hence
| ()
0.8mm(-11))) for all E [n].
((0.0)
j.r _ where $/("") (i [m] m (w (w ..) > ooo /00/V///2), and hence 3.5
(s((
2 क ((-1)n/44 for all j E [+1), r E [m].
Here we take k. and C₁ as 10 and 5 respectively.
Proof of Lemma C.8 We prove Lemma C.8 tions hold naturally by Lemma B.3 and B.4
by induction. When 1' = 0, the fourth and fifth condi-
For the first condition, since we have p that Em [p (0.0) = 0 for all i, k E [n]. So the first condition holds for 1' = 0.
for any j. r, i according to (26). it is straightforward
For the second condition, we have
9i (0.00 x1) Vk f(W(0..) xx)
Fy, (0,0) B
x;) F-y. (W (0.0)
x:) + F (W(0.0) x:) Fy (W/0,,0) x:)
-Be
< Fg. (W(0.) (0.0)
x;) + F ( (0,0)
(0,0)
x;)
( (0.0) B1.
31H4 ) + (P- 1)0 ())(((]) (0.0)
m
(a( my ykM)) + (P 1)o( (0.0) m
(0,0)
mair Ex))
43 1/3 CT,
where the first inequality by F,( xx) 0, the second inequality due to 23 and the third inequality is due to 25
23

---

By Lemma C.6
and the second condition, the third condition can be obtained directly as
(10.0) S exp (yx= /(W Xx) - yı -'( x:)) S exp(Ci).
fk
Now suppose there exists (f. b) & (t, b) such that these five conditions hold for any (0,0) - (t',8) < (f, b). We aim to prove that these conditions also hold for (t', b') = (T,
We first show that, for any 0 S t' < t and 0 - 8/. 82 S b, f(W (f (1), x2) - yk f(W(M xx)
.6) with a small constant approximation error. We can be approximated by m P -P begin by writing out
1/s ".-((We' M)) x:) yh s(w(1)b)) Xx)
=Y: Σ j. F, (M) x:) yk Σ j F,(WW MG) xxs) jE(L1}
jE(L1}
F F-y. ya (W (()
xx) F- p- (W x;) Fg, (Wit (1))) Uk
x;) Fane (W xx)
-h
yr
F F-ve (W
xx)
F- -F y/i
(W!r x;)
m
[o((w (M))))
(t(M)
:84 u)) + (P (( 0:.t [.))
17
(61)b)) year
U/k n.) + (P- - 1)0r(
Ex))
yk.r
F
W Wy
(g))))
Xxk) F- ye
(o(
yerr Me 12) o (( Barr :8k 1
(t' (t)
Ei)) - (P- -1) (9 [P - ()o(((W (1)
(33)
where all the equalities are due the network definition. Then we bound I, Ig and Iz-
For [, we have the following upper bound by Lemma C.4
(W(('.)) (1) x.)
L S | ye (Wl
xx
+ F ve
F W
(C' Mg)
xk)
F m. (W (C' x1)
U2
< 1.
(34)
For [1g]. we have the following upper bound:
(t' M)
(e' (b)
Wk nii}
[2l max
Di 12)),
a
yk.r
m
max
(w.) Werr U: 12)1,
(0,0) yerr Us
T3. .335 SNN
(L)
(.)
32 log log(6n/6)
na
32 log((6//)) na. 8, 3.C'5a,SNRY
max
d
< 0..55,
(35)
24

---

where the second inequality is due to (30), the second inequality is due to the definition of 3 and the third inequality is due Condition 3.1 and (25)
For 13, we have the following upper bound
[PP 1) M1 E.)) (P- 10--10(0w
(C Da.r Ex))
13
m1
[(P - 1) buit
(1)
.i) (P 1)(w
(6) Wk.
EE>T] + 0.25
(eg(6n2/0)
Purira
Py + 101 d na + 0.25
M Poos
Ponorde + 0.5,
(36)
where the first inequality is due Lemma C.5 the second inequality is due to Lemma inequality is due to 5 5 /8) S 1/8 according Condition 3.1
C.3
the third
Similarly, we have the following lower bound
1
(P- 1)0(
E.)) - (P 1)o Mal
(6) Ex))
(1.b.)
(1)
[P - 1)
Ex) (P- - 1)(
Ex>T 0.25
biit
yerr
1og(6n2/8)
(t'
Pue.ri
101
na - 0.25
m
m
(6)
0.5,
(37)
m
where the first inequality is due to Lemma
C.5 the second inequality is due to Lemma C3 the third
inequality is due to 5 5./log (6n2 / < 1/8 according Condition 3.1
By plugging (34) 36 into 33 we have
Vi f (W M
x.)
Uk f(W("
(2) xx) S
[1 + [T2 + Is
m
(1' (t',66)
rrk. + 1.75,
yorrk
m
y4 f(W()M) x.) - yu -f:-ffWWtf
f(W((()))
xx))
[1
[12| + la
(11,66) Por P. -1.75,
which is equivalent to
1 m
y4 f(W((M) x.) yk ff f(W((1.)))
xx)
P
(e(
$ 1.75.
(38)
Therefore, the second condition immediately follows from the first condition.
Then, we prove the first condition holds for (F, b). Recall from Lemma C.1 that
(1,++1 = Pjrr. 1PP 1)2 ((t.)) d'(WW! jrr (.) 113:13 1 = 1i) I.s) Pjra
(1,b)
(t.D)
Bm
for all j E {+1)r r E [m], E [n], (0, 0) S (t) < [T", O]. It follows that
(1.b++1) [(bb+11]
25

---

TTP - 1)2
Σ
p
(1.b)
||5:113 ( E Te,b)
Bm
-5 (etteb) ||Ex [kk E I(ss)
for all i, [n] and 0 S S T H.
If f1,2,
13. 1-1,. then the first statement for (t' (t, b) and for the last (t', <
() are
the same. Otherwise, if 0, we consider two separate cases: Σm= ( p, I 0.9
1.0g
(0110) P ( (6-1.0)
and Dx=
> 0.9k.
(Z-1,0)
m When Σ r=
Pyrrri
I 0.9k, we have
m
(E.0)
(7,0)
Paint
P.k.rck)
p 1,0)
(1-1.01]
112
(15 12) - 1|5.112
Bm
((
( (111))
E:
((--1,0%
1|1E 12
k
pptt (7-00 1.0)
17(P 1)²
| 3/1 ( 1. 11)
(185-1,8(-11)
115:112
Bm
Pyc.r.k
(F-1,00
(TP - 1)2
112:11
< 0.9k + 0.1x
K,
where the first inequality is due to E
(I-1,0
< O: the second inequality is due to |s| 1
m and - ( 1; the third inequality is due to Condition 3.1
(f-10))
(F-1.0)
On the other hand, for when E
> 0.9x, we have from the (38 that
[(W((-1,0
xi) yk [(W(T-1.855))) xx)
1 m
10-1.
")
(8-1,06 -)
1.75
m
1
m
1.0)
1.00
0..1 1.75
0.9k
0.18 0.54k
0.26
(39)
where the second inequality is due to K 10. Thus, according to Lemma C.6 we have
(T- ((
(i-1 u) I exp (yx 1(W0-165- xx) Us f(W(C-1 xx) S exp(-0.26m).
Since s:e S(F-10) C 5 300-1106- we have 5 50-
2 0.80(-1)m according to the fourth
(- 2)
condition. Also we have that 1.
S m. It follows that
S (
exp(-0.26k) 0.8. 0-8+
( (111)
/ ((-1,
|
26

---

According to Lemma B.1 under event Eprelim. we have
||| 12 d. oo 0 (6l))(6n/))). vi E [n].
Note that d = (2(oog (6)) from Condition 3.1 it follows that
|s
Mo.l (0.0011) ) 1|Ex112
|| 12 |sk
k
Then have
2(F.0) Pec
(T-1,0)
1(--1-011 Pyn.rik S n.
P
Paint
which completes the proof of the first hypothesis at iteration (1, = (T, Next, by applying the
approximation in (38). we are ready to verify the second hypothesis at iteration (7, b). In fact, for any (t', (), (t', (4) S (t, b). we have
yi f(W((.))) xx) yk ff f(W(t'))) M2),
1
(01,82)) -P 1.75
xx)
m11
1
(0'-0) Pon.rsk
P
+ 0.1x + 1.75
C1.
where the first inequality is by (38; the last inequality is by induction hypothesis and taking k as 10 and C1 as 5.
And the third hypothesis directly follows by noting that, for any (t, 16) s (T,b),
exp (yk f(w((()) wi) xx) Vi f(w(' M. xx)) I exp(Ci) C2.
For the fourth hypothesis. If (1,2. . 111 1}. then the first statement for (t, (c,6) and for the last (t', 8) < (F, b) are the same. Otherwise, if 6 = 0, according to the gradient descent rule, we have
-
C..) Besr
E.)
(2-10) Werr Ei)
Σ
( o' d'((w ([ ,ypp)) (ye Ed)
Bm
b= V'E
1(P 1)
Σ
H H-1
(1-1bb) d d'(( (F-1) ,x,)) Ex)
Bm
2:.
8' = PEI
H-1
(w (7-100)
Ei)
Σ
(F-166)) . o'((w (T-1)) [BxH))) (ye ps, Edhve
Bm 6'=V'E
yer
v[P (PP-)) 1)
(t 11)
d'((w((@
(7-,5)
E.) 11
Bm
yeer
n(P 1)
(ib') d ((w
(.6')
Ex)) (E., 1(d' +1)
Bm
O0 d'E 1,
=(w(C0) Wear
n(p- - 1)
a 1, 1)
Ei)
Bm
1
'(P-1)
Σ
H-1
El 1116 a' ((
(E-100)
,,)))) (Ee, 11)
Bm
== FEIT-1.'
27

---

H-1
π
Σ b=0 PEEZ
Σ
o d((W yer
yye (ye pa, Ei)
Bm
for any r E S s(f-1,0), where the last equality is by (W
16-1 Ei) > Then we respectively
estimate I,, ls. lg. For 4, according to Lemma B. we have
1.0 " d/2
11 2 10
For 1s. we have following upper bound
H-1
Isl -
100 ((
1 6i.r (xx)) ((E,,,))) #1)
o'((
V= E'E
H-1
1 KEr 1() *
b= V=
H-1
Σ
14 1(-1-)) 209 /dlog(6n218)
S <nc]/e
" 20 v /dl)
where the first inequality is due triangle inequality, the second inequality is due to d'))) E 1), the third inequality is due to Lemma B. the forth inequality is due to the third hypothesis at epoch T-11 1.
For Ig. we have following upper bound
H
Ial S
Σ
14
d'( (T1,6') ..pl))))))))))))))))))))))))))) Kev μ., Et) -
0=
1'E
Σ
V=
nCy 11-1.
||ا|| / 2 (6g((6//)),
where the first inequality is by triangle inequality; the second inequality is due to d'()) E (0, 1): the third inequality is by Lemma B.1 the last inequality is due to the third hypothesis at epoch i 1.
Since d 2 max log(6n2 4C2n| μ lop / 22 2log(6n/8)) we have -(P 1 2
max - 1))44//2 [=]/ and hence- - -PP - 1)14 2 (P 1)14s| + |lol. It follows that
(2.0 w.
Ex) 2 (w (F-1.0)
{.) 0oo, vd/v2.
n:.
for any r E S ((-1,0) Therefore, s/0 C s 1
1,0)
c s(t,0) And it directly follows by Lemma B.3
that (s(000) 2 0.8m@(-1)). 1), Vi E [n].
For the fifth hypothesis, similar to the proof of the fourth hypothesis, we also have
(w((()) (.00 oer
Ei) = (w(2 ((((-10)) E.)
-1)
1,1 1)
11E.11
ye.r
Bm
I
H-1 (/P n(P-1) 1)
((-1,10) a ((w wiet
(w(( (5x, 1('' #1)
Bm
V'
28

---

H-1 Σ V'=0 MEI 1.6'
7
(
d'(w yur yo p1)) (yure, Exhyes
Bm
((--1,U), where the equality holds due to (w (E-1,)) for any i E jjr the same technique used in the proof of the fourth hypothesis, follows that
&.) > 0 and yi j. By applying
(w (2.0 Ju
E 2 (w (7-1.0)
- > 0, j.r
for any i E ((1-1.0) Thus, we have ((0.0) j.r
((t-1,0) ((,))) jor C j,r C 5jr
And it directly follows by Lemma
B.4 that |5 2 (-144)
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) S (t, b) < (f, We aim to prove that they also hold for iteration (f, 6).
Firstly, We prove that 28 exists at iteration (E, i.c., p((.b) 2 -3 10 (log(nn///)///// na for -j,,
any r E [m]. j E }+1} and E [n]] Notice that 2/0r = 0 for j = ys. therefore we only need to
consider the case that j * ys+ We also only need to consider the case of b 6.00 + 1 since (T,b) doesn't I change in other cases according to 21
(T,(1)
When Pirrd
-0.5. 5 /og((n2 / no, by 32 in Lemma C.3 have that (T,6(²)
(T,(())
log(6n² 6m/
1)
jor
E.) S j.ri
+ (P (P_))( (0,0) E.) + 5/
1)
d na < O,
and thus
2 2
(35)
1 P 1)2
(00099) o d(w
(1,69) E.)) 11
frr
Bm
jor
log(6n2/8)
Lirk
3 101
na,
where the last inequality is by induction hypothesis.
(T.60)
When Pine
-0.53 5 /lg (6n /5)/d na, we have
(.b)
(t.bb/))
()P 1))
d' d''(( (( E.) 1|
(1,19))
+
Lir,
Lir
Bm
log (2n6/6//))
T(P - 1)2 30
2 -0.53 5V
na
2Bm
18) na. /
(6n² log(6n2 log
-0.5
log(6n2 1/) d na.
2 - 1
where the first inequality is by E (((.6)
E (-1,0) and ||E:12 < (3/2)ofd by Lemma B.I the second
inequality is due to 5 5 Gn /8)/d na 2 3n d222mm) by Condition 3.1
Next we prove 27 holds for bt, We only need to consider the case of = ye. Consider
((
1 exp(u: [F+1 ( (T.)) x;) F-1 ( (7.6) xx]] (Tb)
(40)
< exp(-y: y: [F (T.)
+1 xx) F-, F_i(W
(7.b) x)))
< exp( y1 x;) + 0.5).
(T
29

---

where the last inequality is by ( (W (7.6) the iterative update rule of Pirr
xx) S 0.5 for j # ys according to Lemma C.4
Now recall
r - 1)²
Pjrú =Para
(((b+)) (1..)) Pjrni
g(bb)) o' (( jrr
(1,b) E.) | 1(i It,b).
Bm
Let (Li,r,s (ina)) be the last time before (T, that (analsan rule from (tyra, barra) (t, b), get
Pjaat
0.5a. Then by iterating the update
(.b)
Pini
(P 1)2
l (llomithenn) 1(w
1( Ei) 2 0) 1(i E I,,6))15.11
(fimni (aami)sann)) jrr
Bm
(41)
(1,b) (P PP-11P -12 glib)) 1((w ((w!l.))
Σ
jr -.) 2 0) 1 L, 1,3)
Bm
We first bound 17 as follows:
||> - (n(P - 1)2 L12/Bm)mmmmmmmm 11:112 S - (nPP-))/Bm)) Bo3///// - 1 S 0.25a:
where the first inequality is by (laraaaaa (-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 a = 4 (T" and T e.
Second, we bound Ig- For (Ij,ra, bana) < (1, (T, b) and yi j. we can lower bound the inner product (w (t,b)
E.) as follows
jor
(t.b) (w Jur
Ex)
2 (w
(0.0) jrr
Eii
1 (t.b)
log(6n2/ /8) na
P-1
d
0.5
0.5
log (og(6n2/8) na
(42)
3 +
a P-1
P-1
P-1
d
0.25
P-1 a,
where the first inequality is by 31 in Lemma C.3 the second inequality is by p/. > 0.5a and (w (0.0) < 1/8 0.1a and / Tog(6n2/8))])
jrr < 2 0P 1) due to the definition of f,,r and B; the last incquality is by
< 0.2a by Condition 3.1
Thus, plugging the lower bounds of (w j.r
(1,b) E.) into Is gives
r)(P 1)2
m
[&
Σ
exp
m Σ (P 1) (1.6)
2rr ..)) + 0.5)
Bm
1(( (t.b) Jrr
. 2 0) 112112
TT n(P 1)2
exp( exp().5) 30²d
Bm
n(P 1)2 1)2
-epl - log (T*) exp(0.5) 30 p
Bm
2
2mm((P1)
Rd
exp(0.5) < S 0.25a,
Bm
where the first inequality is by (40g: the second inequality is by 42 the third inequality is by
= 4lgg (T*); ); the fourth inequality is by Condition 3.1 the last inequality by log(T*) 2 1 and
4'") Plugging the bound of 17, la into 41 completes the proof for p.
30

---

For the upper bound of (29). we prove augmented hypothesis that there exists a i" [n] with (t.0)
Vi = j such that for 1 its T we have that 2ir /Pins- I C. Recall the iterative update rule
of 7irr
(t.b) and Pj.r.i" we have
(p - 1)2
((2bb+1)
(2.0)
e((bb)
o' d'((W) (16) .Ei)) 115i1 1(u 5) A It,o).
Bm
(1bb1)
(1.b) "
E
((( o' d( :/ u))
(t.b)
Bm
((tb)) of ((w jor 11)) 11111/2.
(1.b)
iET,
According the fifth statement of Lemma
C.8 for any 1
s it holds that j yi- and
Ei) (w)))) (1.b)
0 for any (t, (f, b). Thus, we have
(L.0)
(0-10) n(P- 1)? /0-1,000-))
er
||E:- 1
m(P 1)? Bm
(G-Io o
Bm
jr
For the update rule of Ji
according to Lemma
C.8
we have
Σ
Cl o' d'
3ir Ui p))
(1,b)
((2.b) o' d'((w
(1.b)
3: 11)
bc
AE
iET,
(T -1-10 23)1
< C2n
Then, have
(F.0)
(F-00)
C2ne 11,,, (1-1)
Jir
7ir
< max
(1,0)
(1-10)
(f (f-1,6((-1
Pj.r,*"
(P 1)2
02 d/
11-)
1.r
202n E
(43)
maax
p(--100
(P- 1)2o
2C 1C22n||m|11
<
(P 1)
where the last inequality is by F/
(7-1,0)
Piin 2022 2C2n 1)2o Therefore,
(7-1.0)
(2,0)
2029.
(1,0
j.r.a
For iterations other than the starting of en epoch, we have the following upper bound:
(F,b)
(T. (7,0)
4C27.
(2,b)
(2.0)
Pir
Pjr..
(E,b) (T.b)
Thus, by taking C" = 4C2, we have Vir
Pjjrai < C
On the other hand, when (t,
[OR(2T (8)
0), we have
20
(t.b)
log (2T''/1 18)
"
|
1
Fj
20
Bm
12'
where the first inequality is due to update rule of Fimr and the second inequality is due to Condition 3.1
t.b
31

---

When (1, 2
1og(22" (20(2118) /8) 203
,0), According to Lemma B.6 we have
(1.b)
min 2
min
[I ns+ ns-11, IT oe nS ns
Bm
max M 2
Ir ns-1] -1111
7 Bm
(cyca H
B
min gur M) ng max elec10)
ng max ((( |1a1|
i
1.V
IV
0,
where the first inequality is due to the update rule of (1,b) and the third inequality is due to Condition 3.1
7jr the second inequality is due to Lemma B.6
C.2
Decoupling with a Two-stage Analysis
C.2.1 First Stage
Lemma C.9. There exist
Ti = CCn Bm(PP 2d-1 = C.n Cur'BB(P-1))-, 2d
where C3 = (1) a large constant and C₁ = (1) is a small constant, such that
(Ti.0) Pjrt' 2 for any r* E s {r [m] (wh.r En) > 0), E [+13 and E [n]] with yi
naxjrr 35r 0()) for all (t, < (Ti,0).
(t.b)
(t.bh) max (, 0 nn log(()))) for all (t, S (Ti, 0).
max,,r, Ljnn
(1.0) (2 for all 2 T2.
ming, 7j.
maxjrr Pjar, = for all i E [n]-
(T..))
C.2 we have that elc.8) 2 - - / log(Gn2 (8)
Proof of Lemma
C.9
By Proposition
a for all
d
j E (=1), re im], i E [n] and (0,0) I (t. (T", .0. According Lemma B.2 for 3 we have
B = 2 max { (w j.r 1)), (P- - ()
(0.0)
(0,0) j.r
< 2 max ( / 2og ( 2m 20) 0o 111111122210 2 log(12mm/00) 111
O (log(mmm/)) (P 1) vd),
where the last equality is by the first condition of Condition 3.1 Since (t.b)
Link S 0, we have that
max (1,b)) Ljrzs
(L,B) Lj
max J.r.2
35,
log(6n2/08) /8)
< 3 + 1
max { B, (100Vlog(n//))0g log))")))))(())) n/ n\/d// }.
Next, for the growth of (t)
7/.r we have following upper bound
(t,b+))
(1,b)
7)
(t(b))
((w
(1,8) jir 14)) |1laa13
Bm
REI.,
(t,b)
17
||
m
32

---

where the inequality is by |e4 S 1. Note that (0.0) 0 and recursively use the inequality (B + b times we have
(1,b1 -
7 + b)
|1111/2-
(44)
m
Since SNR? =n|1a1//(((P- -1) ) R we have
Ti = -an Bmmmmm (P - 1) 2-= dd
And follows that
15 -
(0)
n(tH +b)
|| -
1111112 C39, MIi mB
m
for all (0,0) S (1, S (T,,
For p recall from 220 that
(pP - 1)2 ((e,) o''( ((wll (1.61)) .Ei)) -1E:l12
(1+1,0)
(t.0) Pyr.
P Puins
Bm
According to Lemma (w (1.6)
C.8 for any r E 0(0000 fr E [m] (wolr E.) > auop /d/ we have
"ir E.) > 0 for all (0,0) s (t,b) S (T*, 0) and hence
-tt100) = (t..0) 7(P 1(P-1) 1)2
(e. 12:12
Byn
For each i, we denote by 11(0 the last time in the period [0, Ti] satisfying that (t.0)
(t.b) m.. S 2. Then for
00 S (t) < (I((0), 0), max,, (1p (tt)) lo(t)) (1) and maxj,r Uir O(). Therefore,
we know that F-1 1 (W(+.) x,), F+1 (W(+)) x,) = 0(). Thus there exists a positive constant C such that - -((r-b) 2 C 2 C2 for S < ((()
Then have
Cnj(P 1)
(1.00) & 2
2Bm
Therefore, ((1.0)
will reach 2 within
Ti Bm - 1)²op'd
iterations for any r' E S(U.0) where C3 can be taken as 4/C.
(t.b) Next, we will discuss the lower bound of the growth of Ygrr For Parii we have
(t.b)
(1.+1) (t,b)
M(t.)) o' o(ww (1.)) Ei) || 1 = (i I,,)
Pirn
Pjrr
Bm
3n(P -1-1 -1)2oud
(1,b) S Pjrri +
2Bm
According to 44) and Pard
(0.0) = 0, it follows that
Pj.r
3r P 1)8oRd((H +b)
(1.8) fjr <
ntHH + +b)
-1144112
(45)
2Bm
(t.b) Therefore, max P Pjrn will be smaller than 1 and (t.b)
smaller than O ([m]| =
O(n (ReSSNRR 0(7) (1) within
To C 1p-1)-207
iterations, where C4 can F- F_((W(t.)) xx), (W((W( x.) = 0(1) in (0,0) < (t,b) - (T2, Thus, there exists
taken as 2/3.
Therefore, we know that
positive constant C such that - (1.b) C' 0 < T2.
33

---

Recall that we denote {i E [na]lus = 3) as Sy, and we have the update rule
11.b+)
(t,0)
"1
quo o' ('(w jr 3/ pu)
(t.b)
30 = -35
Bm
iE9
((b d'( jrr De 14)) 11a11/
(t,b)
For the growth of (t.b)
23 if (w (1,b)
14) 2 0, we have
J.r
(1,b+1)
(t.b)
77
((0
((()) |ell
Bm
iE ns1
iE nsi
2 73.r Bm
(t.b)
1
[C|Z, nS
ns_ ns__ 11111/2
Similarly, if (w)r 11) < 0.
(t.b)
(1,b+1) (t.b)
1 - 1111122
1ir fir +
Bm
Therefore, for 1 E [T2, Til, we have
(1,0) 3ap
Σ
[C n n S [( - II, nS-] |1
Bm
(1,0)(())))
B
(catcs HC Ting)ll
Bm
(c3c41C²
Tin)) 1ll2
Bm
nege, Ctn Ctn||
(46)
8Bm
CCC2 |al2 O( SNR?) 0?)),
(P o)
where the second inequality is due to Lemma B.6 the third inequality is due to < Condition 3.1
C.Crs in 8Cs
And it follows directly from 445) that
3n(P - -1
3C3. (T..0) , Pijmi 0(1),
(T.0)
Pira
2Bm
which completes the proof.
C.2.2 Second Stage
By the signal-noise decomposition, at the end of the first stage, we have
w (2.6) jir = w (0.0)
+
(f.b) |
2 +
(bb
1E "
(3.b)
P
P- i=1
for j [+]] and [m]. By the results get in the first stage, know that at the beginning of this stage, we have the following property holds:
Par > 2 for any r" s(u.0) {r [m]
(Ti.0)
(0.0) Wii
E.) Boo,Vd// (+1) and i e [n]
with 3/4 = j-
max [o(Ti.0)] = max (ax (3,0 log(n/8) log 25r
(T..0) O()) for any {=1)r [m].
34

---

where 9 = n IN.... Now we choose W as follows
Wj. J.r
W (0.0)
20 log (2/c)
Σ1
1 = vi)
E:
P P-1
1
Lemma C.10. Under the same conditions as Theorem 3.2 we have that
02 o "ooppd-ip d- + max ny logloo/////// log(T")/Vd)))
Proof.
[W Wr | -w w(000)
+ | W0000
0 max 35
7j
(11.0)
1 + P 0 ( mm)
max
Σ
(21.0) Pina
&,
(T,.0)
&,
+0 1/2 log(L))) - 1 1 1/22
-(m M
d 1/11 (1 + 1
+ +0 /2n
- 1) d 1/2)
ny =0(m²/3 n SNR. (P - 1) 1/2 m0llmmmmm/mmmm/l//// /m0
+ +0mm 1/2
1/2)
= (2 - + max (3, n
where the first inequality is by triangle inequality, the second inequality and the first equality are by our decomposition of (Z2,,0) W and Lemma B.1 the second equality is by n SNR* = 07)
and SNR ||pll/(P 1),,d//2 the third equality is by n//2. SNR = 0())
Lemma C.11. Under the same conditions as Theorem 3.2 we have that
y: VV (t.b), xx), w 2 log(2/c)
for all (Ti,0) S (t,b) S (T*,0) 0).
Proof of Lemma C.17 Recall that f(w(ta)) = (1/m) Ego (o((w)th)) (0())) + (P
() io
(t,6)
,.))),, thus we have
34 V f( (w()))) xx), W
Σ
(f.b) DiH (I P 1
o' a(ww (1.b) ..i (uEEr,
3.r
3.
= yxr) = (Ev.Ei)
((w arr
(L,b)
16
m=1
(T.11, Exr
((ww.r (0b) (.14))200og(2/ [ij yir)
|| 112
1
(t.b) ((w
00.4) jw (0.0)
p
Σ
o' d'((
(t.b) (Et) (u: (0,0)
j.r
jor
m j.r
m
J.r
1
Σ
d(lw
(t.b) (.))20ll 1 ΣΣ o' ((W 2.) log(()/)
(t.b)
1(Ev.
| 12
j=ye.
J.r
Kürp Er)
o' d'((w (t.b) Di14) )20 log(2/e)
Σ
('(( (t.b) j.r
Ex
1=1
j.
35

---

1 Σ
o' ((w
(1.0) -.)) log(2/c) Σ
(' (6,b) for [1)2000og((//00mm//ggn)/)))
3=
3.
I10
1 m
o' (( (1.b) (n SNR dt
1
Σ o
((w)
(t.b) M,H))), (47)
j.r
where the first inequality by Lemma B.2 and the last inequality is by Lemma bound each term in (7)) respectively.
B.1 Then, we will
For 11o, L11. 112,
114, have that
Hoo &
[ul s (nn/log(n/O)) SNR d 1/),,
(48)
[12| (0))
For j = yı and r E 5 according to Lemma
C.3
we have
log(4n2 15) 5m
(P - 1)((
2 (P- ()(w))))))
jrr
(0.0)
Ei) + Pj.r.
(1.b)
d
a
log (4nn² (44m2/))_
> 2 8 5m
d
1.5 8 0,
where the first inequality is by Lemma C.3 the second inequality is by 5n log [og(4n2 18)
< 0.5: and
the last inequality is by 8 1.5. Theretore, for 1g, according the fourth statement of Proposition C.8
we have
Ig ((tbb) (20log(//)) 2 22og(2/c).
(49)
By plugging 48 and (49
into
(47) and according to triangle inequality we have
W V ((W((MMM)
W 19 1oo [nl
[ua 2 log(2/c), [12
which completes the proof.
Lemma C.12. Under Assumption 3.1 for (0. S (t, < (0), the following result holds.
IV (w(*))) >1> s Ommx(]] (P 1) of)))) (. (W(s.)).
Proof. We first prove that
(f( x, x;)ll= = Omax([]] (P 1)o,
(50)
Without loss of generality, we suppose that D: Then we have that
| IV (W(()))
xx)lE= s m
Σ
lo' (ww
(t.b) p)) + (P 1)o'(( (1.b) E
j.r
Σ m 1.
d'((w(!
(1.b) 11))1|1112
P-1 Σ
o' d(w(
(1,b) (Wrr))
1.)))
m
3.
< <max 12, 2(P (LP-1)apval
where the first and second inequalities are by triangle inequality, the third inequality by Lemma B.I Then upper bound the gradient norm (VLs(W(1bb)) as:
2
VLZ, (W((.))) <
Σ I
((uf(W((1.)) xx)) (V xx)lFF
B
36

---

Σ
O(max (1 (P 1)o, 8 (vi((W()))) xxi)))
B
EEI
0 (max (P- 1)- B
Σ VEL,,
-l' (Wii((Wii))) xx1)
S (max [ (P () (W((.))
where the first inequality is by triangle inequality, the second inequality is by 50 the third inequality by Cauchy-Schwards inequality and the last inequality due to the property of the cross entropy loss- -l' < l.
L
Lemma C.13. Under the same conditions as Theorem 3.2 we have that
[| w ||
- W 1| 2 nLI,, (Wi(..))
27c myc
for all (Tı,0) S (t, - (T" 0).
Proof of Lemma C.I3
We have
[W w
|| (((,).)) W -2 (W(Ld)), W(())) W - 7/2 V LS (W((M)) 1|?
2n
Σ
e (L.b)
Ui (WIM xx) (V((W( x1), W 17 LZ.,, ( (w(sd))
VETu.
2n
((. ( Wi x.) - -log(2/e)) - 1 Vss ( ((
2
B
TEI
27
(()))(((()))) xx)) c/22 VVIV "L (W(++)))?
EEL,, III (W(+b)) nJc,
where the first inequality is by Lemma C.11 the second inequality is due to the convexity of the cross-entropy function; the last inequality is due Lemma C.12
Lemma C.14. Under the same conditions as Theorem 3.2 with probability 5, we have that
(zu.s(((tt.)) (W(+.))- Tzza, (zzaor (W(101) Sc,
for all (Ti,0) S (1, S (T*,))
Proof.
LIreess) (W(1.)))
Liin (Wi(0))
((nr(W(a.)) xx:)) (W(t))
<
B
VEI..b
1
Σ
(i ,x,) (. ,2))
<
B
V
w (1.) 14) + (P - -1) (w/(())))) w (1.))
<
j.r 1E3
B
m
TET.1.
Hn 1P)
Vulll2pp//2og((0 (lgg(//))))))) + Hn(P - 1)2 -12 20 / ddlog(6n2
<
Bm
Bm
56.
where the first inequality is due triangle inequality, the second inequality is due to | S 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.I and the fifth inequality is due to Condition 3.1 D
37

---

Lemma C.15. Under the same conditions Theorem
3.2 for all Ti S t S T", we have lo max (0(Voog(mn/)) oo - 1)op d), O / log.////////// n/0) log(T")/ //d))////))))))))))))) /d) Be-
maxj,r,
sides,
LI. (W(L)) | - W P
E
T +c
(s Hi)H
for all Tı <t T Therefore, can find iterate with training loss smaller than 2e within T = Ti+ |W(Z1)- - - (1/ee Ti + O(n 2mddd iterations.
Proof of Lemma C.15 Note that jaxj j.rie = max {o (Ooooommn//)) oo(P
1) p V dd, (n logg//////////o///////
can be proved in the same way Lemma C.9
For any T S 7*, by taking summation of the inequality in Lemma C.13 and dividing (s - H on both sides, we obtain that
1
Σ
II (W(r.)))))) <
[ -W 12
te.
T
17(8 T.H
(L.)
According to the definition of T, we have
1
E
Lu (W(t.)) $ 2c.
(T - TT)H (T.
Then there exists an epoch Ti S S T such that
1 H-1
Σ
Lz,, (W(((:)) 2Ec.
H
6=
Thus, according to Lemma C.14
we have
Ls(W(t.0)) S 3c.
Lemma C.16. Under the same conditions as Theorem 3.2 we have
Pinalny = O(SNR-* 2) Σ (bb) (1.b)
(51)
for all j.3 [+1},, r. [m] and (T2, 0) bt) < (T*
Proof. Now suppose that there exists (0,0) < (T,0) < (T*,0) 0) such that Σ
Ei Pjrz //g.."
(t.0)
((SNR=2) for all (0,0) < (t,0) < (T,0). Then for Pjr according to Lemma C.1 have
(t,b)
(fr =Gj
(1+1.0) (1,b)
(t,b)
d'l(w
(t,b9))
D 14)
Bm
bc -ES
Σ
gle.b
a'(lw!!
(1,65))
.0t 14)) 11al
Jir
LES nde
Pjsú
(+1.00)
-p j.re
(P 1)2
o'(lw
(1.99))
(.)) 11:112 1(y: = j),
Bm
follows that
(T,0)
juna
38

---

Σ
(T jri
=
(f-100) n(P -))
( 1
(F-1662-1)
0 d((w
Bm
T 1,0)
1/P (1)-2)
Σ
(T-1 11)
E
Bm
SES 1 1-
jrr
7/P T
(T-1) jan +
2
min
(52)
8m
1=1
where the last equality is by the definition of -fT inequality is by Lemma B.I and the fifth statement of Lemma C And
s(T-1) as fi [m]] y/
(T-1)
J.
{w
j.r 1.: >
0}; the last
(T.0)
(T-100)
1.6
(T 1T.1
Di 12) |1a113
,'
Bm
j',"
iES+
< 3'
(T-1,0)
Hnll
max iES+
1 11)
(53)
14 L
According to the third statement of Lemma
C.8 we have maxies_
leg 1-10 n, 1e ( 1. Then by combining
C2 min
52 and
53
we have
(.,) E jire
E Pirn
11,00
(P- 1)
min
OSSNR 3)
(54)
(
(T-1.0)
16C₂ 1ull3
Σe (2,0) Piri (1.0)
On the other hand, we will now show
OSNNR 2) for 2 T2 by induction. By Lemma
B.I and
52
we have
n
3n(P - 1)3ozdn
(T2.)
(T2-10) joni
2Bm
3mp( P
2Bm
And, by 46 we know that at t = T2, we have
ne304 CT2n
(T2.0)
8Bm
Thus,
(T2.0) -=11 Paas
I OSSNR 2).
(T2.0)
S ((T00) (2.0)
Suppose
OSNR 3), According to the decomposition, we have:
(T.b)
(:144
= (w j.r
(0,0)
Tip))
7i. T
(T)))
+
1
(T.b)
116:1122(E...
(E. Dipl) +
Σ (T,b) Aind
1 D.44). (55)
p
p 1
i=1
39

---

And we have that
(0.0) W
+
1
Pirg ||E, 22 (E:, T1
(T,6)
+ P-1 1=1
|16 12 (S: 074) (T.b) j
P-1
1=1
n - 8/2 +
41121112 2 v log (6n/6)
P
(T)))
Jins
Jnd(P - 1)
1=1
O(SNR "330
(T.6)
< 8/222 +
vd
(T.0) - )rr
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 (w)) (T.b)
suppose (w (T,6) Jor (sspe) is persistent through out the epoch. Then, without loss of generality,
ji pe) > 0. Thus, the update rule of y is:
(1,
7i.
(1.0) (t,b)
7
(T.b))
Σ
((TT)b)
|el12
= 7,g
Bm
iETT.onS+nS)
iETT,nss_nn ns.
2 1i
(T,b)
"
min
e! (T.b)
[T n nSi iEZT
max
[T n 1- nS- |1/41/2.
Bm
iETT,
Therefore,
(T+1,0)
(T
7
SS nsil max (T,
(T IS_ ns_1I] 111113.
min
(56)
Bm
And, by
52
we have
(T(P - 1)202 1) (-1) (T.0)
(++10)
((T,6
Pira
Pira
max
(57)
8m
Thus, combining
56 and
57
we have
(T+100)
11=1
Pind
(T+1.0)
7ir
Σ n
(T.0)
(P- (1-1)203dh o=-1))pdn max 1e (T,)))))))
iii Pj.ra
< max
(T.0)
8 min of (T,
15 nil - max J.T.D. (T)) )
|s_ns= ||
$ OSNNR 2),
(58)
where the last inequality is due to the induction hypothesis, the third statement of Lemma Lemma B.5 Thus, by induction, we have for all Ti S < T that
C.8 and
Σ 1 Pira
(1.0)
OOSNR ?).
(2.0)
And for (71,0) < (1, S (T", 0), we can bound the ratio as follows:
En
(c.b) P
(c.0) Pira
@SNN 2),,
(6.b)
(2.0)
where the first inequality due to the update rule of p and (t,b) proof.
(t,b)
p Thus, we have completed the
D
40

---

C.3 Test Error Analysis
In this section, present and prove the exact upper bound and lower bound of test 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 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 C1 such that when n| Ci 1y for time t defined in Lemmal C.15 we have the test error
P P(xx))D(V # sign(f xx)) s exp (opppp
where C2 = 0(1.
Proof. The proof similar the proof Theorem E.1 in Kou et al. (2023). The only difference is substituting & in their proof with (P - 1)
C.3.2 Test Error Lower Bound
In this part, we prove the lower bound the test error in Theorem
3.2 We give two key Lemmas.
Lemma C.18. For (Ti,0) S b1 < (T*,0), ,), denote g(E) Σ j(P j(P exists a fixed vector with [|v|]2 0.060, such that
((w .E)) There (1,b)
Σ fe(L)
Io(I'E + g('' 2 4C
max JE(+1}
Σ 1 (1,b)
(59)
for all & E Rdd
Proof of Lemma C.18
The proof is similar to the proof of Lemma 5.8 in Kou et al. (203)). The only
difference is substituting & in their proof with (P 1)
Lemma C.19 (Proposition 2.1 in Devroye et al. 2018)). The TV distance between N(0, aa and (Vv La) is smaller than l|v|2//20p-
Then, we can prove the lower bound of the test error.
Theorem C.20 (Third part of Theorem 3.2h. Suppose that rellpai12 < Csd(P- 1yo, then we have that L (W(+))) 2 p +pp 0.1, where C3 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 & in their proof with (P-
D Proofs for SAM
D.1 Noise Memorization Prevention
The following lemma shows the update rule of the neural network
Lemma D.1. We denote ((t.b) L'(v.. f(W(..) x,)), then the adversarial point of ((t.b) is W(t.)) + 2(1,) where
(1.8) T EiETa EpE(p) e(u,b)/ j. y.o'(( (1,b) x..)) )xi
m
VV LI, (W(tb))
Then the training update rule of the parameter is w (1+1.))
w (1.b)
Σ Σ
((e.b)
(( +
(1.0)
(((bb) xi Xix
j.r
Bm
SEI.
PE|P
41

---

W ji
(1,b)
E
Σ Ci(i)
d' ((
(1.b) jo xi.p) (Cr.j.r xig))s xix
Bm
&
PE[PI
(1,b)
E
e((cb) d' d'( .up)
(t.6)
Bm
1 - 1)
((tbh) o (( ((w))) .E) + (2 .
Bm
NoiseTTrrm
We will show that the noise term will be small if we train with the SAM algorithm. We consider the first stage where & S Ti where Ti = mB/(11nn) p (2). Then, the following property holds.
Proposition D.2. Under Assumption
3.7 for 0 S & S Tı, we have that
(0.0)
(0.0) (0.00) 0,
7ir
P j.na E.rg
(60)
0 S 1i
(t.b)
S 1/12,
(61)
0 S P jir
(t.b)
S 1/12,
(62)
log Dog(Gn2// (ō) n.
0 2 (t.bb)
2 -3 -100 d
(63)
Jir
(T..0) = 2(1). Besides, Yj.
Lemma D.3. Under Assumption r E [m]. j E [+1] and i [n].
3.1 suppose
27
28
and
29
hold at iteration t. Then, for all
(ww(t.))
w
3.r
(0,0)
14) - Y5
(t.b)
I SNR
32 log (6n/8)
d
na.
(64)
Wir .Ei
(0.0)
(ft.b)
5
log(6n2 /8)
nax + V/s
(65)
P 1
P
d
(w Jin
(0.0) w jir
{x)
1
(t.b)
5
log(6n2 /8)
d na, y/a
(66)
P
P
Proof of Lemma D.3 Notice that 1/12 ot, if the condition (61).
63 holds, 27
28 and
(29h also holds. Therefore, by Lemma C33 we know that Lemma D.3 also holds.
also
Lemma D.4. Under Assumption
3.1
suppose 61
62
63 hold at iteration t, b. Then, for all
j {=1} and i [n], F,
(W (t.6)
0.5. Therefore
0.3
E 2 0077.
Proof. Notice that 1/12 < a, if the condition 61 62
63 holds,
28 and 29 also holds.
Therefore, by Lemma C.4 we know that for all + y and E [m]. F3
x;) S 0.5. Next we
will show that for = Vi. F (1.)),) x;) S 0.5 also holds.
According to Lemma D.3 we have
F2 (1.
1
o ((w (4.6) m.m) + (P ()w (9 E.))
x:)
< 2 max
(1
(1,6) M. (P 1)l (1)
j.r: 1
32 log (6n/
(0) j.rt D. (P- 1)]( jort E.), SNR
< 6 max
no.
d
log(6n2/0
5V
na, 7ij
(1,8)
p
(t.0) J.r.
d
< <6max { 3, /
32 log 2))
log (6n2 (n2/6)
d na, 1/ (1.b)) (t.b)
d
na, 5/
Pirsi
42

---

S 0.5,
where the second inequality is by (64) (65
and (66); the third inequality is due to the definition of 33;
the last inequality is by (55
662
Since FW (L.b)
x.)
[0, 0.5] we know that
0.3 1 exp(0.5)
2 242 1 expl -0.5) 2 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
(w (t.b) E 2 0, E Itb and j yk. then
(t.b)
Ek² < O.
Proof. We first prove that there for & s Tí, there exists a constant C2 such that
VV Www.,,,,w/(t))) S ( C2Pooy
Recall that
1 Σ (l..( x.)), B
LI,, (W.(.b)
REI
we have
V Vwg, LI (WiL
Σ
Vwy. (W.(e,) x,))
VEIL.
1
Σ
yell (ur (urf(W(t))) xx)) (W(t.)),
B
TET,,b
Di (o' ((w)))
(6.b) pd) 1 + o' ((w (2b) Ex) (P))-(P-1)) 1))]]-
Bm
iEI..b
We have
VW,, (w(1.b))122
S Bm
(((t))) of d''((w
(1)
11
+
(((b))
d'(( ((w)).))) (.)) .P (P- 1)
Bm
iEZ,,
2
iEZz,
2
< 0.7mm
||1l|12 + 1.4(P 1)m
opVd/BB
< 2Pm famm-topp opvd/BB
and
(w(())) =>
Σ
V Vw,,, Lz,
(w(1b)) Vw 2m(2Pm top/ d/B
2.r
leading to
Vw [VWLz,(W((0)]] = Bm
From Lemma D.I
we have
(2((.)) jur
Ex)
T
Ww
LI, (W((
Σ
4 yio yio'((w S0 xip) Ex)
mB
JEI. PE(P]
VW (W (W((.)))))>
(1.b) jy: (P - 1)0 ((w (/ (1.0)
(.) Ex)
mB
VEL
(10) jyk (P- 1)o0 ((w
(Ex) Ex)
43

---

Σ (t.bl1 o' ((w (t.b)
Vi Ex)
O.8B(P Vlog +. 20g((6n
mC2P P Bd
PP.15(P 1)
TopV mVB B
(67)
where we the last equality is by choosing T by (66) we have that
mvB Cs CoPoyyva
Now we give an upper bound of (w Ex)
llo (og(6n2//1 16) na,
(w 10 jir Ex) & max
(0) J.r E,)1, 5V
(tab) Jinz
s 11/(4(P- 1)).
(68)
Combining (67 and (68) completes the proof.
Lemma D.6. Under Assumption 3.1 suppose 61 holds for t, +,
63
hold iteration b. Then
also
Proof. Now consider the SAM algorithm. Recall that
- - - 1)2 ((t.))) o'( jor + f E.) ||3:13 1(y; = 1(1 I.M.)
(t.b+1) Pjr.a
Bm
Case 1: & Tz,b- In this case, clearly we have that (1,b+1] = p/.r - 1/12.
Case 2ti E IG, and (w (w!( Ei) 2 0, then by Lemma D.5 we have that (w j.r + <(t)
Ejrr Ex) < O.
therefore we have that Pinn
(1,++1) (1.b)
Pina I 1/12.
Case 3:i I and (w (t.0) can not reach a constant order,
j.r Ei) S then by (66) and triangle inequality, we can conclude that Pj.r.i
log (6n2
T S (P
(1,b) w (0.0) ()
j.
Ei)l + =5V
no.
Then we can give an upper bound for ((+11.)) since we only take one small step further,
11,b+) 1) Pir.
< (P -
1)
(t.b)
S 5V Wir
(0.0)
log log(6n2/0)
+
T[P -P-1))²
Bm 2do < 1/12.
d
Proof of Proposition D.2 We will use induction to give the proof. The results are obvious hold at
t = 0 as all the coefficients are zero. Suppose that there exists T < T₃ such that the results in
Proposition D.2 hold for all time (0,0) < (t, b) < (T - 1, - 1). We aim to prove that 61). (62).
also hold for iteration (T 1,6).
First, we prove that (61) holds for hold for iteration ( - 1, b). Notice that (t.0) ((w(1)))
3r
(t.b+1)
7jr
(t,b)
Bm
(t.b)
a' d'lw
j.r 3/ p1)= |1 <
(t.b)
+ m
7 ||
VEI..
where the last inequality is by the fact that 10 ((1bb+2)) < 1 and o' < 1. Notice that 7 - < Ti, we conclude that,
( Dr
<Ti
eTi(nn/B
11 || < 1/12.
m
44

---

Second, by Lemma D.6 we know that (62) holds for (T - 1,
Last, we need to prove that 63 holds (T 1, The proof is similar to previous proof without SAM.
(F-1,b-11 1) < -0.5(P- 1)8- 6v
/ooeoooe /ogg/n2/08) (6) d na, by 31 we have
When Bi.ck
(0.0) (r (f .Ek) < (w
16--1 1)
log(6n2 /0)
jr
[x) + P 1
+ p P-VV
d
na
log (og(6m2/// 10)
$ P
d na.
and we have
(2
(1 E.)
V W Li,
(W(T-13-1)) 1F
Σ
( 1 g y
mB
1 1,8-1 PE[P]
o' d'((w) --115-11
Jir xx. .xx.)))( Ex)
V
-1811 (W(T-1
E
/ 1, 11 yi
mB
iEI7 1.5 1
(P-
1)00 (1-1 (.)))EE. Ex) + +1/ jyk.
(P jyk
(lw' for (E)) Ex)
frr
( 11 j. o' d'((w) 1--13-1) .yit4) (p, Exi
+
1,r
1
O.8B(P- - 1)02 dlog (6n (n2//) /0) + 0.4 /ttp:/llllllllll(@/r11 18)
mC2Pop Bd
TV Bo log(6n2//)
< C4
m C4 By log (6)
CsPVd
1
log(6n2/0) na.
$ PV
d
and thus (w) (-515511 1)
efft -1-11,511 1) E.) < 0 which leads to
jrr
(1-11b-1) + PP ((F-116611
oo((
(--1551)
.e)) -1
= 2i,,
Bm
12y; = -)) 1( E I -1ū-1
11111b-1)
j.r,
Therefore, we have
(- 1-1,6-11
> -PP 1)3 - 5P
log(6n2/0
d na.
When p! 1- 2 -0.5(P - 11)8- - -5/ log na. we have that
fjr,i
(T,),
2 (-11b611
+ n(P - 1)2
((T-115-11
115:113
j.r.ā
j.zi
Bm
> 1-11b611
0.4m(P P 1)2
2do p
Bm
log(6n22/
2 - 1),3 5P
na.
45

---

Therefore, the induction is completed, and thus Proposition D.2 holds.
(t) Next, we will prove that jrr can achieve (2(1) after Ti mB/ x iterations. By
Lemma [36] we know that there exists C3 Ti epochs such that at least H batches in these epochs, satisfy
B 3B
S-_ ng nI
for both = +1 and y/ = -1. For SAM, we have the following update rule for (t.b).
7jrr
33.r = =Yjo (1.6)
(1.b+1)
Σ ((bb))
o' d'
(t,b)
+
(t.bb)
jr :bs pJ) ||
Bm
VETU,nSS
+
"
(((.))) d(
(1.0)
+ 2(t.))
Dr 11) |
Bm iEI,,nS
3rr
If (w (1.b) (..) (1.b)
.14) 2 0, have
(1,b41)
= (t.b) (t.b)
p(())
|
Bm
iE nn
2 7j.r Bm
(t.b)
(0.3)
n 0 Si| 0.7 S+ n ns- 11
If (w
(1.b) (..) (1.b))
114) < 0, we have
F3r = -17 + 17
(1.+1)
(2,b) (t.b)
Σ
((())
|
Bm
n8 ns
iE 08
2 Vis +
(t,b)
1
(0.3)Ia, 0 S+ 0
S-1 0.7 0.7 n 0 nSli) -111
Bm
Therefore, have
(T,.0)
2
" (0.3 cTTI
H.H 0.25B - 0.7T:ng)1
Bm
"
0.
Bm
2
cacaTim cacsTin||u113
16Bm
C3C1 2(1).
192
Lemma D.7. Suppose Condition 3.1 holds. Then we have that [ [| = -e(oova d) and
(T..0)
ww (T..0)
fir .jm) = 22(1)
<w (Ti.0)
-j.r .jm) = -((1),
= 2 max[1(w) (Ti.0) 8,J. jrr" pu)l (P () (T..0)
Besides, for (t.,)) and (/1,b) defined in Lemma B.3 and B.4 have that
|S(T1.0) -Ω(m), Vi [n]
|s ( = I(n), vj E (=1). [m].
46

---

D.2
Test Error Analysis
Proof of Theorem
4.1 Recall that
W
(t.b) 3.
w
(0.0) ju
+
(1.b) 7,5
-ut + -1111112
1
(t.b)
|16:112" E..
P
by triangle inequality we have
w J.r 1|2 - // |/
(21.0111
10000001
11a
-
1/
(1,0)
-1111112 11 + P-1 1
15:l 5
i=1
1
||1/
+
12(P 1) (o2d/22)-122
vn
12
By the condition on og and Lemma B.2 we have
(T) o
112 ([[l"] - l12) A vd).
By taking the inner product with مr and E., we can get
(w)) pe) (w]]r)) (0.0) 14) + -2. P
(1.b)
(e.b)
1
12
(1,b)
11:1ll22 2. (E..
Pars
and
(w(h (1.6) Jrr EE) = (w (0.0) (0.0)) E.) + "Jrr" |||11122 -p, +
(t.b)
Piras" -11.v111"" 61211 : (E),
(s.b)
p -
y==
(w/000 (0.0)
E.) + 3or |||111222 2. (p. +
(t.b)
(t:b)
Phink
p
+ p
Par |1 112 (Ev,
(1,b)
Then, we have
(w(T (T..)
j.r
.Jp) = (w/ .pp) + +13.1
(0,0)
(T.,0)
+
P
1
Pira
(T.0))
1| (E..j))
2 13s - 1(
(T1.0)
(0.0)
(Ti.0)
|| 11)
Pinnt
P
75 1ogg o 0ul|ua12
(Ti.0)
n 12(P - 1)
22log(6n//
(T..) fr
and
(w (w (T),0) -j. -114) = (w (0.0)
:j.r: j44 - -3-3
(T.00)
(T.00)
11..1122 (E.
P -
joci
1 n
< - -Ji
(Ti.0)
-jor | +l(w(0.))
(0.0) +I(w
1p -3r
(7:.0)
||E 2 (EE. 13)
< T-j
(Ti.0)
/ (2mm/))
To 0ollull2
1111200p / 22log(6n/0)
12(P-1)
(T)
< 2 15.
47

---

(2).0) where the last inequality is by the condition on og and Tirr = 92(1). Thus, it follows that
(T..0) (w jr juu = (2()), (w -jir .i4) = -I()).
(T,.0)
By triangle inequality, we have
kw (T..0) j.r
-10) S
((w
j.r
(0.0)
.1))|
h3
(T.
+
P-1
1
Σ pooel 11: I 14)
2
1
+
11
(o2d/2) 1 -11el1200 (lgo((6n/8)
12
P-1 12
2 +
+
n1
|Ma2 (log((n/00/((op))
12
6(P
1)
S 6'
and
((w.7..0) (T,.0) .) 5,r -Ei $ ((w/0r) (0.0) - + 1/ (T.00) ||11112 KK,, 6:) + p 1
1
lp (T1.0)
+ p - 1 (
1p (2,00)
||E ((E,,E))
1
I 2 8 + |112 pp' v (2ngg(6n/5) + 12(P 1)
12
+ P2 - 1)
n
(0 (0,d/2) 2) 2220 v /dlog(6n2///
1
1
$ 28 +
+ 6
f11u1l2 op vlog(6n/0)
12 - 1)
< 6'
This leads to
3 = 2 max
(T) [.J. jor ,1)1 (P - (w (Ti,0) , = 0(1).
In addition, we also have for < Ti and j = yı that
(w (1.6)
E2) - (w (0.0)
j.r <
3rr
1 (e.M) (t.b)
1
P Ppir T/rr
|||11122 (K
Ip -11E::112 2
p
2 (1.b)
|1|111222 114,
1
Σl
11E,. E.)1 |15.r 2
1 1
it''"
2 12 |1ml 1144,
n1
115:1152 1(E2, E.)
12(P - 1)
2 12
||1| oop (2log(6n/5)
12(P 1) (dgd/2) (0112))-2206 220 /dlog(6n2/0)
n
|111 o, v log
3(P - 1) v /og
n
12
2 - | pp- Vlog(Gn/0).
Now let 5(0.0) denote fr : (w oar < >
ood,Vd)}} and let 5/0.0) denote fi E [m] y: =
J. (w birr {,) > ooonVd) By the condition on do. we have for t < Tj that
(w (1,b) 30r
Ei) 2
[w (0,0)
Jrr Ei)
22
48

---

for any E 5 or E 5,r Therefore, we have 5: C S (1) and 5/000 C s(f1.)) and hence
0.8 (- | 1,000 | (s() = 9(m),
0. (2250(- S |1S),r | ((n),
where (()) is the CDF of the standard normal distribution.
Now can give proof of Theorem 4
Proof of Theorem After the training process of SAM after Ti, we get WTT,0), To differentiate the SAM process and SGD process. We use W to denote the trajectory obtained by SAM in the we have that
proof. i... W(T,,0) By Proposition D.2
(T1.0) (0.0) + 2/
+j-3,
(Ti.0)
fe
(T,,0) Pira
(T.0)
E:
P
[
2
i=1
(69)
(TT.0) where (TT.0) = ((), Pjra E [0.1/12], Lyra E [-3 - (Og/lo/(mn//// (ogs6622/8)/dnn 0j. Then the 1ir
(Ti.0)
SGD start at W(..0) W(T.o), Notice that by Lemma D.7 we know that the initial weights
of SGD (i.e., the end weight of SAM) W(0,0) still satisfies the conditions for Subsection C. and
Therefore, following the same analysis in Subsection C.1 and C.2
we have that there exist
t "in='c-'mmd" mmd-mmndd---- 222 such that Lg (W(r:))) < c. Besides,
n
((1,0) (0,0) jir
(w)) j.r
+ 7i.
(1,0)
(1.0) J.r
(t.0) Jist
P
1
P
E
for j E [+1] and E [m]] where
(1,0) Yir
OSSNR3
Σ
Piran
(1.0)
(t.0) Pjrz E [0,
(11.0) a-a,0]
(70)
ig|m]
Next, we will evaluate the test error for W(+.0), Notice that we use (f) as the shorthand notation of (1, 0). For the sake of convenience, we use (x, D. y) 'D to denote the following: data point (x, y) follows distribution D defined in Definition 2.1
and y is its true label. We can write out the test error
as
P(x & sign( x)))
= PPxxN)DDUU/(C (W/(W(()) x) 0)
=P(xNIDD (y/ (W(() x) S 0, y + 7) + P (V x) 50, = 7) = p P( (7() x) 2 0) +(1 p) P (O x) S 0)
(71)
S p Sp+ P (î (! x) ().
where in the second equation we used the definition of D in Definition 2.1 It therefore suffices to provide an upper bound for P P(xx (f((W()) x) < 0). To achieve this. we write x = (ūp, E) and get
Of (f(w())
x)
1
Σ
Ti (w)) 3. (1:) + (( (2 E))
m
j.r
1
o ((w Tir:
Tn)) + (P- 1) (1) E)
m
Σ
( TMm) + (P (w D.r E))
(1)
(72)
m
49

---

The inner product with j D can be bounded as
(w.. (1ee = (w: (w. U.rt Tre) + N.r +
Bos 111:1122 (E,, 01L)
(P
+ (P 1)
11::112 (E:. D1)
2 logg(bn//)
(w (0)
Dirt 71L) + 7/.
(2)
P-1
op FpllA1l2 (o/d/2) -1
D.r,i
+
Σ
(t)
1=1
(w Birt T1 + /gr
(1)
0 Vlog(n/0) (Popd( O(SNR 3)
(1)
'.g.r
(w Dir: Dm)) + 11-0(Vlo((n//)) Pop/ Pop/|lee[22]]76 (8)
(w. "ir Tn)) + ((7 (t)
Q(1),
(73)
where the inequality is by Lemma
B.I the second equality obtained by plugging in the coefficient
orders we summarized at (70); the third equality is by the condition SNR = |2 Papl:/ the fourth equality is due to || 2 2 2 log(n/6) in Condition 3.1 for sufficiently large constant C' the equality holds; the last equality is by Lemma D.7 Moreover, we can deduce in similar manner that
(w (1) H.rt 012) (w(
F,r
T1)
-Y-g.r B.r +
>
B... 11E.112 -pp)
+ (0)
1=1 B.r. ||E:1122 (Ei, 0P))
S (w (0)
7,r U11) --ior
(()
(74)
splllall2 (ozd/2) + / 2 (0g(66/8) .
(1) Bin
Biri
1=1
1=1
(w (0)
7.r TT) -(7")
= 2(1) 0,
where the second equality holds based on similar analyses as in (73).
Denote (({) Σ, a(w (t)
g.r E)) According to Theorem 5.2.2 in Vershynin 2018), we know that
for any 2 it holds that
c2
P(g(E) Eg(E) 2 2) S exp
(75)
1eg Lip
where e is a constant. To calculate the Lipschitz norm, we have
19(5) 9(5'))
Σ 0 c((w (1)
T.r. 5) Σ @(ww (1)
B.r 5'))
o (w T.r. {)) - o(( B.r. E2>)
Iw (c)
J... E E'>l
w (t)
1.rll2 ll ||g E'1l2.
50

---

where the first inequality is by triangle inequality, the second inequality is by the property of ReLU; and the last inequality by Cauchy-Schwarn inequality. Therefore, we have
|| 2.rll2
(76)
and since (w E) NO, || [/ (0)
-r. (..111))) we can get
g.r1100p
Bg(E) =
Bo (w p.rt E)
Op
(t)
g.r
122
2m
2
Next seek upper bound the 2-norm of w (t) tion, namely:
j First, tackle the noise section in the decomposi-
T1
Σ
i0 Piaaa
1|E:112
E:
1
Σ
( 2 jrn
11E:112 2 +
Σ
Prrssppp 122 112 .E;, Ei,)
(1)
1 Si
n
dd < 40p p
Piin
(t)
+2
Pinn Piinia
(t)
(160,' 2) (20 V dloo((6n2/0 /8))
15h Sn
=07 d'
(1)
3/2 /log(6n2/6) + +.20, 3/2
[
2
(0) 2
Plintl
P
i=1
O(o," 26,20"1
hnd + =000 2d 3/2)
(1) 2
Piin
(0)
10 'n'''nn n
+ + (o 2d 2d13
Σ
L0(e) ).r.al
1=1
s 50(0,2d n
where for the first inequality, we used Lemma B.1 for the second inequality, we used the definition of p. pt for the second to last equation, we plugged in coefficient orders. We can thus upper bound
the 2-norm of (4)
J.r as:
|[ fo <
ww 10 + 7/rr
||
+ p
Σ
( Pank 11:12 E
2
n |1 + O(P op dd 1/2 -1/2) Σ
< wlw!0 [w]fll2 12 + 3/.
(t)
jini
1/2 1/2)
!1)
(77)
janj
where the first inequality is due to the triangle inequality, and the equality is due to the following comparisons:
7/ 1114
1(P-1 on 12 122 ||1111 SSNR3
O(P-lop d 1/2 1/2). Σt= P
=O(P-1 07 1/2 1/2 (1(1122)
1/2
= 0(1)
51

---

based on the coefficient order Σ
(1) 10 - @(SNR the definition SNR = / vil) Di=13 Pizz/ 3,r
and the condition for d in Condition 3.1 and also |w 1. ((ouv based on Lemma D.7 With this and 773) we analyze the key component in (80):
Σ, ((((
(1) g.r 311))
(1)
(P - 1)0p Σm
|w frrll
op d- d-1//2n-112 1/2). E. Pjirs
((1)
(78)
IV W
min 1) 222
It directly follows that
[o((w
(P- p
w Dr lla > 0.
(79)
2n
Now using the method in 75 with the results above, we plug
74 into
72 and then (7 to obtain
P(x.)))> (f
(W
x) 5))
S <P(DD
Σ Cal(
!1)
-ir {) 2 (1/(P- 1)) Σ o ((w 8.r 0T1)))
= -P(xoo)~D
g(E) Eg({) 2 (1/(P-1)) Σ ( a((w D.rt D.r: F(ee)
Jp
|| w
(2)
g.rllz
2T
r=1
(11((-11))
- (1(1--1)) E, o((w (1)
m)1)) (op/ v 2
m
1/
(
2
G.rt
S exp
of E
(t)
E, ((w Fir T14))
27
1/
exp
(P- 1)0p Em
(w( -1,rl12
Σ, o T.rt TT1)
exp(c/27) exp
0.5c
(80)
(P- 1)o, Mn
where the second inequality is by (79) and plugging 776) into fact that (s t)² 2 /2 t2, Vs, f 0.
75). the third incquality is due to the
And we can get from and 80) that
Σ, S(w((w (2)
2
P P(x.D.)> (TI(((1))))))))) x) S 0) S explc/2m)exx
exp
0.5c
T.rt
(P- - 1)0p Em ww
-B.r 11
- Cminfo 2d Pož -ta -1=
exp
2x
exp
0.50 0.50'imn min(o
2d- Pot -taaaa 2] 2
where C 0(). the last inequality holds since of S 0.5Cd~logg(1/ and d >
20-1p 1-10,p log(1/e).
52


Reviews:
[{"review_id": "muZzHjgBjL", "reviewer": "Reviewer_TY4G", "summary": "This paper presents two theoretical contributions regarding benign/malign overfitting of two-layer convolutional ReLU neural networks. For an idealized data distribution, it i) gives the conditions (with respect to the dimension of the data and to the signal strength) under which benign/malign overfitting occurs when training the network with mini-batch-SGD; ii) shows that training the network with SAM leads benign overfitting more often than SGD and in particular in many situations where SGD suffers from malign overfitting. These findings are illustrated by a numerical experiment conducted in the idealized situation used for the theoretical analysis.", "questions": "1) Equations Line 175 and 176 are given as a definition (Definition 3.3), but appear to be a result of the data distribution and of the network architecture chosen. Why are they presented as a definition?\n2) The numerical experiment is performed with gradient descent instead of SGD, while the theoretical analysis deeply relies on SGD. Why? Are results identical with SGD?\n3) Similarly, are the theoretical results for SAM the same if optimization is performed with gradient descent?", "limitations": "See above regarding mathematical assumptions. Societal impact is not addressed.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 4, "strengths": "The two contributions presented in this paper are significant because Contribution i) improves on the literature, in particular with respect to the closest work (Kou et al., 2023), which proposes exactly the same analyses but for a gradient descent training; Contribution ii) is, up to my knowledge, the first such result for SAM. The latter result provides a comprehension on the well-known fact that SAM helps networks to generalize better than SGD. This is a real breakthrough, which should be welcomed by the community.", "weaknesses": "I was very enthusiastic by the results presented in this paper but I became disillusioned when reading it. The presentation is disordered and in particular, notations are not all defined, making the results (even informally stated) very difficult to understand. After having a glance at the appendix (which seems to be rigorously written), it seems that the manuscript is an assembly of results extracted from the appendix, unfortunately awkwardly built. As for me, it is a pity because I think that the results are important and of interest for the community, but the presentation is inadequate and not clear enough.\nComments:\n1) The abstract misses to state Contribution i) (which constitutes a large part of the manuscript and which is clearly reminded in the conclusion).\n2) Definition 2.1 is difficult to understand. In particular, the parts “signal contained in each data point” and “the other are given by” are unclear. Besides, even though this model seems common in the recent literature, a discussion regarding its limitations, for instance independence between labels and covariates, would be appreciated. The same remark can be done regarding the architecture of the neural network considered (Section 2.2). In addition, Condition 3.1 is discussed only quickly Line 142 ; a deeper discussion regarding the role of the variances and the polynomial degrees appearing would be appreciated.\n3) There is no link between statements/results in the paper and their counterparts in the appendix. Since the reader is supposed to juggle the two, it would be helpful to know where are the formal statements and the proofs in the appendix.\n4) I understand that authors try to explain the derivation of their main results but I am not very comfortable with informal statements, first because of the notation problem previously stated and also because some conclusions are given with vague explanations (in my case they are difficult to grasp) while they seem to be quite difficult to obtain formally. This is the case for instance Line 218 regarding the symmetry of $\\rho$.\n5) Line 117, the authors invoke the discontinuity of the gradient to justify the need of an analysis based on an other technique than the Taylor expansion. As for me an even better argument is that the ReLU function is formally not differentiable everywhere.\n6) The statement of Theorem 4.1 is a bit unclear: the parts “we train/we can train” could be replaced by passive forms, “neural networks” refers, as far as I understand, to the chosen architecture. In addition, mentioning SGD in a result concerning SAM may throw the reader. It could be specified beforehand that SAM training uses intrinsically SGD (but at a point which is not the current iterate).\n7) Lemma 4.3 is true only for a particular choice of $\\tau$, as stated in Theorem 4.1 (I am not sure that this is clear in Lemma C.5). Since this result is important, this should appear in Lemma 4.3 and discussed after.\n8) The related work section (Section 6) appears at the end of the manuscript but is an enumeration of papers. Such an enumeration is generally well placed just after the introduction. Placing the related work section after the result statements makes sense if it discusses technical differences with the closest papers. As it happens, this work seems to be based to a certain extent on (Kou et al., 2023). Thus, it would be enriching to discuss the contribution, and particularly the technical novelties (GD → SGD/SAM), of this work with respect to (Kou et al., 2023).\n9) Mathematical remarks: Line 47, $t$, $\\mathbf W^{(t)}$, $\\boldsymbol \\mu$, $L_{\\mathcal D}^{0-1}$, $p$, $\\Omega$ are not defined; Lines 47, 150, 240 and after, the expression “converges to $\\epsilon$” should be replaced by “converges to $0$”; Line 66, $l_2$ → $\\ell_2$, Line 68, absolute values in $|a_k/b_k|$ seem useless; “omit logarithmic terms” should be defined explicitly; Line 97, $\\mathbf W$ is not defined; Line 100, “is a collection of” means “is a matrix”, doesn’t it? Line 101, $[n]$ is not defined; Line 113, it is specified that $\\sigma_0^2$ is the variance of the normal distribution but not Line 80; in Equation (5), the 2-norm should be a Frobienus norm; what is the utility of Equation (6) with respect to the equation Line 175? Line 121, $S_i^{(t, b)}$ → $\\tilde S_i^{(t, b)}$; Line 219, $T^*$ is not defined; Lines 218 and 226, the authors could remind what are $H$ and $B$; Line 226, I understand that $|\\cdot|$ is the cardinality but this is not stated.\n10) Typographical remarks: Line 9, “for the certain” → “for a certain”; Line 27, “with minimal gradient” → “with minimal gradient norm”; Figure 1, blue and yellow are inverted; Line 102, “cross-entropy loss function” and “logistic loss” are redundant; Line 115, “indicator function” → “the indicator function”; Line 156, “Bayesian optimal risk” → “Bayes risk”; full points are missing Lines 202, 207, 214; Line 237: “iteration” → “iterations”; Line 251, full point instead of colon; Figure 2, y-label is cut; Line 299, “an generalization” → “a generalization”.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback. Due to space limits, we will address your major comments and questions as follows. We will revise the corresponding parts accordingly as well as address the minor points in the final version.\n\n**Q1**. The abstract misses to state Contribution i).\n\n**A1**. Thank you for your suggestion. We will add the contribution i) in the abstract of our final version. \n\n---\n\n**Q2**. Definition 2.1 is difficult to understand.  A discussion regarding its limitations, for instance independence between labels and covariates, would be appreciated. \n\n**A2**. Thank you for your feedback. We will add more explanation of Definition 2.1 and the architecture of neural networks in our final version for a clear presentation. \nRegarding the “independence between label and covariates,” we believe there is a misunderstanding. For any covariate $x=[x^{(1)},... x^{(P)}]$, there is exactly one $x^{(j)} = y\\cdot\\mu$, and the others are random Gaussian vectors.  For example, the covariates $x$ could be $[y\\cdot\\mu, \\xi, \\ldots, \\xi]$, $[\\xi,\\ldots, y\\cdot\\mu, \\xi]$ or $[\\xi, \\ldots, \\xi, y\\cdot\\mu]$. The signal patch $y\\cdot\\mu$ can appear at any position. So $x$ and $y$ are not independent. We will make it clearer in the final version.\n\n---\n\n**Q3**. It would be helpful to know where the formal statements and the proofs are in the appendix.\n\n**A3**. Thank you for your suggestion. We will add pointers in our final version. In particular, Theorem 1.1 is the informal statement of Theorem 3.2 and Theorem 4.1. Lemma 3.4 is the informal statement of Lemma B.8. Lemma 3.5 is the informal statement of Lemma A.6. Lemma 4.3 is the informal statement of Lemma C.5.\n\n---\n\n**Q4**. Not very comfortable with informal statements, for instance Line 218 regarding the symmetry of $\\rho$.\n\n**A4**. We’re sorry for the confusion. We utilized the term \"symmetry\" to describe a situation where the summation of $\\bar{\\rho}$ corresponding to different samples yields similar values. More precisely, the difference between these values can be bounded by a small constant: $\\sum_{r=1}^{m}\\zeta_{y_i,r,i}^{(t,b_1)}-\\sum_{r=1}^{m}\\zeta_{y_k,r,k}^{(t,b_2)}\\leq\\kappa$, as indicated in lines 203-204. We will remove the word  'symmetry' from the paper and revise this part accordingly.\n\n\n---\n\n**Q5**. Lemma 4.3 is true only for a particular choice of $\\tau$, as stated in Theorem 4.1.\n\n**A5**. We’re sorry for missing the condition in Lemma 4.3. It requires $\\tau =  \\Theta\\Big(\\frac{m\\sqrt{B}}{P\\sigma_{p}\\sqrt{d}}\\Big)$. \n\n---\n\n\n**Q6**. Discuss the technical novelties of this work with respect to (Kou et al., 2023).\n\n**A6**. The key technical challenges and novelties of our work compared with Kou et al. 2023 are highlighted as follows:\n\n- We studied SGD rather than GD. As the mini-batch update of SGD only utilizes only a small part of samples, different samples contribute to the update of coefficients differently. Consequently, the noisy samples may deviate the learning process. Thus, we have developed techniques to control the update of coefficients at both batch-level and epoch-level. Specifically, Lemma B.7 bounds the batch-level update, and Lemma B.8 controls the epoch-level update and aligns the update of batch-level and epoch-level.\n\n- We also provide a novel analysis for SAM, which is very different from GD/SGD. SAM has a completely different neuron activation pattern from SGD. The activation pattern is based on the perturbed weight $\\mathbf{w} + \\mathbf{\\epsilon}$  in SAM, rather than the ​​unperturbed weight $\\mathbf{w}$ as in SGD. The perturbation $ \\mathbf{\\epsilon}$ introduces difficulties in the analysis. We discussed the difficulties and our techniques to tackle them in Section 4.1. More specifically, we decompose SAM updates into SGD step and perturbation step and connect them through Lemma 4.3. We show that if a neuron is activated by noises in the SGD step, it will subsequently become deactivated for the perturbation step. This technique is new and has never been used in prior works such as Kou et al., 2023. It is pivotal for the analysis of SAM.  \n\n---\n\n**Q7**. Mathematical and typographical remarks.\n\n**A7**.\n- Line 100, The phrase \"a collection of\" can be interpreted as a new matrix or as a tensor of the weight matrix.\n\n- Equation (6) is a variant of Line 175, which we refer to for technical clarity. By further 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. \n\n- $T^*$ is defined in section B.1.2 as $T^* = \\eta^{-1} \\text{poly}(\\epsilon^{-1}, d, n,m)$. \n\n\n---\n\n**Q8**. Why Equations Line 175 and 176 are given as a definition?\n\n**A8**. Thank you for your suggestion. We will present it as a lemma instead of a definition in the revision to ensure clarity. \n\n---\n\n**Q9**. The numerical experiment is performed with gradient descent instead of SGD, while the theoretical analysis deeply relies on SGD. Why? Are results identical with SGD?\n\n**A9**. Our results also apply to GD, since GD can be viewed as a special case of SGD where the batch size is equal to the dataset size ($B = n$). We have additional experiments of SGD (batch=1024) on the real data set in Appendix D, and it is also consistent with our theoretical analysis. We have also added an experiment of SGD with mini-batch size 10 on the synthetic data in Figure 1 of the uploaded PDF. If we compare Figure 1 of the uploaded PDF and Figure 2 of our main paper, the comparison results remain the same. \n\n---\n\n**Q10**. Are the theoretical results for SAM the same if optimization is performed with gradient descent?\n\n**A10**. Yes, since SAM with full gradient is a special case of SAM with stochastic gradient when the batch size is equal to the dataset size (B = n), our theoretical results also hold for SAM with full gradient.\n"}, {"review_id": "KMHGjalgza", "reviewer": "Reviewer_jpeN", "summary": "The paper aims to provide a theoretical basis for the superiority of SAM over SGD. Different from former explanation based on Hessian information, the authors firstly discuss the loss landscape of non-smooth neural networks like two-layer convolutional ReLU networks. Notably, the paper proves that under conditions of harmful overfitting using SGD, SAM can deliver benign overfitting and outperform SGD in terms of generalization error. The experiments are conducted to demonstrate that SAM can outperforms SGD in terms of generalization error by mitigating noise learning and enhancing the efficiency of weak feature learning.", "questions": "See weakness.", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The motivation is clear and novel. Previous researches focused on shallow model with implicit smooth loss. As a more challenging setting, the paper studies two-layer convolutional networks with ReLU units. \n\n- The theoretical analysis is convincing. In terms of benign overfitting, the authors discuss the phenomenon and condition in SGD firstly. In harmful overfitting, the authors theoretically prove that SAM can outperform SGD in the early stage of learning.", "weaknesses": "- The experiment results may not be reliable. In Section 5, the author sets the learning rate to 0.1. However, according to recent research [1] and my practical experience, I suspect the learning rate is too high and may lead the model to overfit easily. The model may not easily overfit if we use SGD with suitable learning rate. Hence, it might be better to add an experiment and discuss the influence of learning rate.\n\n- Potential overclaiming based on the analysis and experiment. At the beginning of the paper, the authors claim that SAM can help model learn weak features effectively. However, I do not see obvious related content that can support the authors claim. I recommend the authors modify their claims or try to discuss the relationship between SAM and weak features more clearly.\n\n[1] Andriushchenko et al. “A Modern Look at the Relationship between Sharpness and Generalization”, ICML, 2023.\n\n-------\n\nAfter reading the responses from the authors, my concerns has been well addressed. I will keep my score as weak accept.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback. We provide our responses below and will make the corresponding adjustments in the final version.\n\n**Q1**. I suspect the learning rate 0.1 is too high … It might be better to add an experiment and discuss the influence of the learning rate.\n\n**A1**. Thank you for drawing attention to the recent work [1]. We'll cite and discuss it in the related work section. Additionally, we appreciate your keen observation regarding the learning rate. We indeed used a learning rate of 0.01, not 0.1, in Section 5. We will fix this typo in the revision.\n\nFurthermore, we've conducted an ablation study on the learning rate, which is presented in the uploaded pdf. Specifically, we experimented with learning rates of 0.001, 0.01, 0.1, and 1 under the conditions described in Section 5. Figure 2 in the uploaded PDF shows that for learning rates of 0.01 and 0.001, the patterns of harmful and benign overfitting are quite similar and consistent with our Theorem 3.2. Since SGD with a small learning rate needs a longer time to converge, we didn't finish the experiments for extremely smaller learning rates due to the time limit. But given that Theorem 3.2 holds when the learning rate $\\eta$ is sufficiently small, we believe that the phase transition for a smaller learning rate would also align with the results from learning rates of 0.01 and 0.001 if training longer. \n\nInterestingly, we observed that larger learning rates like 0.1 and 1 enhance SGD's generalization performance, but are still worse than SAM. The benign overfitting region expands for learning rates of 0.1 and 1 when compared with 0.01 and 0.001. This observation might be tied to the phenomenon highlighted by Reviewer CEVG. We'll incorporate these findings and further discuss the impact of learning rate variations in our revision.\n\n[1] Andriushchenko et al. “A Modern Look at the Relationship between Sharpness and Generalization”, ICML, 2023.\n\n---\n\n**Q2**. I recommend the authors modify their claims or try to discuss the relationship between SAM and weak features more clearly.\n\n**A2**. Thank you for your feedback. We believe there may be some confusion caused by the term \"weak feature\". In our paper, 'weak features' refer to features when the signal-to-noise ratio is low. We will clarify this point by removing “weak” in our revision. Our claim was based on the fact that, for benign overfitting, SGD requires the norm of features to be at least in the order of $d^{1/4}\\sigma_p$, while SAM only necessitates the norm of features to be almost a constant. This weaker requirement of SAM allows it to learn features more effectively compared to SGD. Figure 1 in our manuscript aims to provide an intuitive visualization of this result. \n\n\n\n"}, {"review_id": "0zea9tMFhD", "reviewer": "Reviewer_CEVG", "summary": "This paper studies the question of why SAM generalizes better than SGD on a specific binary classification task. The task looks like a special case of the sparse coding model where the relevant parameter is a signal-to-noise ratio (SNR). The authors present theoretical results on the performance of SAM and SGD for this specific task. More specifically, their main results say that SGD requires higher SNR in the data distribution to generalize, while SAM can generalize even with smaller SNR.", "questions": "- The binary classification considered in this work seems to be an instantiation of the well known sparse coding problem. It would be helpful to motivate the model by presenting connection to the well known sparse coding problem.\n\n- Connection to large learning rates GD/SGD. Recently there have been several works trying to understand why large learning rate GD/SGD generalize better than small learning rate ones. It would be clarifying to see if that's also the case for the model considered in this work. It seems that there are recent works that show that large learning rates help with generalization for a sparse coding problem (e.g. [1]). Given this, it would be a nice clarification to test at least empirically if one can achieve SAM's performance with larger learning rate GD/SGD.\n\n- Related to the comment about the large learning rate, Q. Is the result for SGD also true for Gradient Flow (continuous dynamics of SGD)? \n\n- The theoretical results show that SAM updates help prevent memorizing the spurious feature (or harmful overfitting) during the early stage of the training. For the model in this paper, does that mean you can switch from SAM to SGD after a few epochs and basically achieve the same performance as SAM?\n\n- The required perturbation radius ($\\tau$) in theory scales inversely with $\\sqrt{d}$, which seems much smaller than what people choose in practice. Does this requirement consistent with the experiments in the paper? I think it's an important point since a sufficiently large $\\tau$ distinguishes SAM from SGD.\n\n- The original motivation for SAM was to encourage the landscape around the solution to be more flat. Does that intuition also hold true for the sparse coding problem considered in this work? Does the prevention of harmful memorization lead to a flatter landscape when you consider the flatness around the iterates of SAM?\n\n- Does the oscillation of SAM (dynamics characterized in the previous works mentioned in this paper) also occur in the training dynamics of SAM for the model in this paper? If so, does the oscillation another explanation as to why SAM does not learn spurious features?\n\n- In the experiments, it seems that only full-batch GD/SAM are tested for the synthetic model setting. Could you comment on the behavior for SGD and Stochastic SAM? \n\n- Also, many recent works have studied un-normalized variant of SAM (usually referred to as USAM), where the ascent step normalization is removed. Are the main results for SAM also true for USAM or is the normalization necessary for the main results?\n\n- In the related work section, it would be nice to see discussion on other techniques (large learning rates, label noise etc) as there have been many works on their effects on generalization (e.g. [1] [2] [3]).\n\n\n\n- [1] Learning threshold neurons via the \"edge of stability\" (https://arxiv.org/abs/2212.07469) \n- [2] SGD with Large Step Sizes Learns Sparse Features (https://arxiv.org/abs/2210.05337)\n- [3] What Happens after SGD Reaches Zero Loss? --A Mathematical Framework (https://arxiv.org/abs/2110.06914)\n- [4] Label Noise SGD Provably Prefers Flat Global Minimizers (https://arxiv.org/abs/2106.06530)", "limitations": "Please see the comments in the Question section.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The paper has a very comprehensive end-to-end analysis for neural network learning, which is usually very technically complicated.\n- The paper is overall well-written and it was easy to follow.\n- Figure 1 and Figure 2 have very nice resemblance.", "weaknesses": "- Discussion of the results seems not quite comprehensive; see the comments below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback! Due to space limits, we answer your major comments and questions as follows.\n\n**Q1**. It would be helpful to motivate the model by presenting a connection to the well known sparse coding problem.\n\n**A1**. Thank you for bringing our attention to the sparse coding problem. This is definitely a good way to motivate our study. We will emphasize the connection of our data model to the sparse coding problem in the revision and add the related work.\n\n---\n\n**Q2**. Connection to large learning rates GD/SGD … It would be a nice clarification to test at least empirically if one can achieve SAM's performance with larger learning rate GD/SGD.\n\n**A2**. Thank you for pointing out the relevant studies and suggesting to test larger learning rates for GD/SGD. We will cite these works and study the effects of learning rate variations in our revision.\nFirstly, we apologize for the oversight in Section 5, where we mistakenly mentioned a learning rate of 0.1, which should be 0.01. This will be corrected in the revision.\n\nAdditionally, we've conducted a detailed ablation study on the learning rate in the uploaded PDF. Specifically, we tested the implications of larger learning rates  0.1 and 1 with the same condition in Section 5. Figure 2 in the uploaded pdf indicates that slightly larger learning rates indeed boosts the generalization performance of SGD. The benign overfitting region is enlarged for learning rates of 0.1 and 1 when contrasted with 0.01. This trend resonates with the findings of the studies you've recommended. 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 larger learning rate exhibits improved generalization, it still falls short of matching SAM's performance.\nWe'll integrate these observations and provide a more comprehensive discussion on learning rate in the revision.\n\n---\n\n**Q3**. Is the result for SGD also true for Gradient Flow (continuous dynamics of SGD)?\n\n**A3**. We require only an upper bound for the learning rate $\\eta$. Therefore, our theory is applicable for any sufficiently small learning rate and should potentially be extended to gradient flow.\n\n---\n\n**Q4**. Does that mean you can switch from SAM to SGD after a few epochs and basically achieve the same performance as SAM?\n\n**A4**. Yes, under the model studied in this paper, transitioning from SAM to SGD after certain epochs yields performance comparable to that of SAM. This observation aligns with practical findings, such as Figure 9 in Andriushchenko and Flammarion (2022) .\n\n---\n\n**Q5**. The required perturbation radius in theory scales inversely with $\\sqrt{d}$, which seems much smaller than what people choose in practice. \n\n**A5**. The required perturbation radius in Theorem 4.1 is $\\tau = \\Theta(m\\sqrt{B}/P\\sigma_p\\sqrt{d})$, which depends on many factors and also hides the constant. Take the synthetic data as an example, we have $(m \\sqrt{B} / P \\sigma_{p} \\sqrt{d} = 10  * \\sqrt{20} / (2 * 1 * \\sqrt{10000}) \\approx 0.22 $ which is not small at all. We choose $\\tau = 0.03$ due to the constant omitted in the bound.\n\n---\n\n**Q6**. Does the prevention of harmful memorization lead to a flatter landscape when you consider the flatness around the iterates of SAM?\n\n**A6**. Previous papers focus on the flatness (sharpness) to explain the success of SAM. Our paper provides a different perspective for why SAM generalizes better than SGD. In our current analysis, we were not able to show that the prevention of harmful memorization can make SAM converge to flat minima. Yet this is a very interesting and important question to study in the future.\n\n---\n\n**Q7**. Does the oscillation of SAM also occur in the training dynamics of SAM for the model in this paper?\n\n**A7**. Thank you for this insightful question. Our current theory does not imply the oscillation of SAM dynamics. In our experiments, we did not visualize the dynamics of SAM iterates due to the very complex landscape of the ReLU network (2-layer ReLU on synthetic data and ResNet50/WRN-16-8 on CIFAR10). It is also not easy to visualize the neural network weights in a 2D/3D space. We plan to investigate this question in our future work.\n\n---\n\n**Q8**. Could you comment on the behavior for SGD and Stochastic SAM?\n\n**A8**. The comparison keeps the same for SGD and SAM with smaller batch sizes. We have provided additional experiments for them. Please see Figure 1 in the uploaded pdf. \n\n---\n\n**Q9**. Are the main results for SAM also true for USAM or is the normalization necessary for the main results?\n\n**A9**. The main results for our results relied on the normalization and we carefully characterize it in our proof. Our results do not directly apply to the unnormalized SAM.\n\n---\n\n**Q10**. In the related work section, it would be nice to see discussion on other techniques.\n\n**A10**. Thank you for your suggestion. We will discuss these techniques and related works in our revision.\n\n---\n\n[1] Ahn et al., \"Learning threshold neurons via the\" edge of stability\".\" arXiv preprint, 2022.\n\n[2] Andriushchenko et al. \"Sgd with large step sizes learns sparse features.\" ICML, 2023.\n\n[3] Li et al., \"What Happens after SGD Reaches Zero Loss?--A Mathematical Framework.\" arXiv preprint, 2021.\n\n[4] Damian et al., \"Label noise sgd provably prefers flat global minimizers.\" NeurIPS, 2021.\n"}, {"review_id": "7XhEXiEvMQ", "reviewer": "Reviewer_EhwP", "summary": "This paper presents an in-depth theoretical examination of Sharpness-Aware Minimization (SAM) in the context of feature learning. The authors identify the challenge of overfitting in the training of these networks, a problem that becomes more prominent as the model size increases. The traditional gradient-based methods like Gradient Descent (GD) and Stochastic Gradient Descent (SGD) are identified to suffer from unstable training and harmful overfitting. The authors highlight SAM as a promising alternative that has demonstrated improved generalization even in situations where label noise is present. However, they point out the current lack of understanding on why SAM outperforms SGD, especially in nonlinear neural networks and classification tasks.\n\nThe core contribution of this paper is filling this knowledge gap by providing theoretical reasoning on why SAM has better generalization than SGD, particularly in the context of two-layer convolutional ReLU networks. The authors characterize the conditions under which benign overfitting can occur when using SGD, and further demonstrate a phase transition phenomenon unique to SGD. They formally prove that under these same conditions, SAM can achieve benign overfitting, and therefore has a distinct advantage in terms of generalization error. Specifically, the paper highlights the ability of SAM to mitigate noise learning, which prevents it from succumbing to harmful overfitting early on in the training process. This aspect of SAM allows it to facilitate more effective learning of weak features. The theoretical findings are supported by experiments on both synthetic and real data, bolstering the credibility of the presented theory.", "questions": "Could you please share your thoughts on the potential behavior of SAM within the gray region depicted in Figure 1?\n", "limitations": "I believe the authors could provide a more extensive evaluation of potential limitations inherent in their research. Engaging in a comprehensive discussion regarding these constraints would strengthen the overall rigor and transparency of the study.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- **Addressing a Significant Issue**: The authors tackle the important issue of overfitting in the training of large neural networks. They delve into a key challenge facing the field of deep learning, thus making their work relevant and timely.\n\n- **Theoretical Contributions**: This paper provides a strong theoretical analysis of why Sharpness-Aware Minimization (SAM) outperforms Stochastic Gradient Descent (SGD), particularly for two-layer convolutional ReLU networks. This contributes to the current understanding of how to improve the generalization of large neural networks.\n\n- **Comprehensive Study**: The authors carry out an in-depth comparison of SAM and SGD, using both synthetic and real data. This robust approach adds to the validity and comprehensiveness of their findings.\n\n- **Novelty**: The authors claim that this is the first benign overfitting result for a neural network trained with mini-batch SGD. This claim, if validated, could signify a novel contribution to the field.\n\n- **Clarity and Organization**: The paper appears to be well-structured and clearly written, with a good overview of the background literature, a clear statement of the problem, and a detailed account of the authors' contributions. \n", "weaknesses": "- **Limited Scope**: The paper's focus on two-layer convolutional ReLU networks with fixed second layer might limit the generalizability of its findings to other types of networks. Further research could be needed to determine if SAM has similar benefits for different architectures.\n\n- **Presentation Could Be Enhanced**: The current version of the paper could benefit from further polishing in terms of presentation. For instance, the y-axis in Figure 2 lacks clarity — it's not immediately evident whether it represents a normal scale or a logarithmic scale. Also, the 'dimension' it not clear. The term 'clear samples' could be more accurately represented as 'clean samples.' Additionally, there's inconsistency in the usage of 'P-1' and 'P', which should ideally be standardized throughout the paper for improved readability and comprehension.\n\n- **Partial Theoretical Results**: The authors provide a thorough analysis of both benign and harmful overfitting scenarios for SGD. However, the corresponding harmful overfitting regime for SAM appears to be absent from the theoretical results, which could make the comparison somewhat incomplete.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your strong support. Below, we provide answers to your comments and questions, and we will ensure that the revisions are made in the final draft.\n\n**Q1**. Further research could be needed to determine if SAM has similar benefits for different architectures.\n\n**A1**. Thank you for your constructive feedback. Since this is the first paper towards formally understanding why SAM outperforms SGD for ReLU networks, we choose to avoid introducing unnecessary complexities in the data/CNN models and the analysis to make the result easy to follow. This ensures clarity and ease of understanding. Exploring more complex data/CNN models certainly presents an exciting avenue for future research, and we plan to investigate it in subsequent studies.\n\n---\n\n**Q2**. The y-axis in Figure 2 lacks clarity. There's inconsistency in the usage of 'P-1' and 'P'.\n\n**A2**. Thank you for your valuable feedback, the y-axis represents a normal scale with a range of 1000-21000. We will address the issues mentioned above in our final version. As for the usage of ‘P-1’ and ‘P’, the data has P patches, where P-1 patches among them are noise. Therefore, ‘P-1’ and ‘P’ both exist in our paper. However, we understand this may cause confusion and will find a way for clear expression in our final version.\n\n---\n\n**Q3**. The corresponding harmful overfitting regime for SAM appears to be absent from the theoretical results, which could make the comparison somewhat incomplete.\n\n**A3**. The main focus of our paper is not to provide a complete characterization of benign/harmful regions of SAM but to show that SAM has a larger benign overfitting region than SGD (Figure 1). This suffices to explain why SAM can generalize better than SGD. We leave the investigation of harmful overfitting regime for SAM as a future work.\n\n---\n\n**Q4**. Could you please share your thoughts on the potential behavior of SAM within the gray region depicted in Figure 1?\n\n**A4**. We conjecture that SAM generalizes badly on the gray region due to our empirical results in Figure 2 where the blue color represents high test error. It is an interesting future direction to prove the phase transition between harmful overfitting and benign overfitting for the SAM algorithm.\n"}]

Author rebuttal:

We want to thank all the reviewers for their valuable comments. In the uploaded pdf, we include additional empirical results to address reviewers' concern including Figures 1 and 2. 

+ **Figure1**:  To address Reviewer TY4G’s concern that our synthetic experiment is only performed with gradient descent instead of SGD,  we have added an experiment of SGD with mini batch size 10 on the synthetic data. If we compare this Figure 1 of the uploaded PDF and Figure 2 of our main paper, the comparison results should remain the same and both support our main Theorems 3.2 and 4.1.

+ **Figure2**: Reviewers jpeN and CEVG raise an intriguing direction that whether tuning ​​learning rate can get better generalization performance and achieve SAM's performance.  Therefore, We've conducted an extended study on the learning rate. Specifically, we experimented with learning rates of 0.001, 0.01, 0.1, and 1 under the same conditions described in Section 5. The results show that for all learning rates, the patterns of harmful and benign overfitting are quite similar and consistent with our Theorem 3.2. We observed that larger learning rates such as 0.1 and 1 can enhance SGD's generalization performance. This observation might be related to the phenomenon pointed out by Reviewer CEVG.  Importantly, for all learning rates, the benign overfitting regions remain 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. 


Meta-review: The paper tackles an important question of understanding SAM vs SGD. The reviewers are all positive about the paper and authors response was well received. We encourage the authors to clean up the presentation, connect better with the appendix, provide concrete definitions and formal results alongside high level intuition and pay careful attention to typos. Please also include the discussion and experiments on the effect of learning rate.
Decision: Accept (poster)