Title: Why Does Sharpness-Aware Minimization Generalize Better Than SGD?
Venue: NeurIPS 2023
Accepted: yes
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?

Abstract

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

1 Introduction

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

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

Foret et al. (2020) provided a PAC-Bayes bound on the generalization error of SAM to show that it will generalize well, while the bound only holds for the infeasible average-direction perturbation instead of

∥

μ

∥

2

Dimension

d

Benign

Overfitting

Harmful

Overfitting

∥

μ

∥

2 =

Θ

(

d

1/4

)

Signal Strength

∥

μ

∥

2

Dimension

d

Benign

Overfitting

∥

μ

∥

2

= ˜

Ω

(1)

SGD

SAM

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

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

Theorem 1.1 (Informal statement of Theorems 3.2 and 4.1) . Let p be the strength of the label flipping noise. For any ϵ > 0 , under certain regularity conditions, with high probability, there exists 0 ≤ t ≤ T such that the training loss converges, i.e., L S ( W ( t ) ) ≤ ϵ . Besides,

For SGD , when the signal strength ∥ µ ∥ 2 ≥ Ω( d 1 / 4 ) , we have L 0 -1 D ( W ( t ) ) ≤ p + ϵ . When the signal strength ∥ µ ∥ 2 ≤ O ( d 1 / 4 ) , we have L 0 -1 D ( W ( t ) ) ≥ p +0 . 1 .

For SAM , provided the signal strength ∥ µ ∥ 2 ≥ ˜ Ω(1) , we have L 0 -1 D ( W ( t ) ) ≤ p + ϵ .

Our contributions are summarized as follows:

We discuss how the loss landscape of two-layer convolutional ReLU networks is different from the smooth loss landscape and thus the current explanation for the success of SAM based on the Hessian information is insufficient for neural networks.

To understand the limit of SGD, we precisely characterize the conditions under which benign overfitting can occur in training two-layer convolutional ReLU networks with SGD. To the best of our knowledge, this is the first benign overfitting result for neural network trained with mini-batch SGD. We also prove a phase transition phenomenon for SGD, which is illustrated in Figure 1.

Under the conditions when SGD leads to harmful overfitting, we formally prove that SAM can achieve benign overfitting. Consequently, we establish a rigorous theoretical distinction between SAM and SGD, demonstrating that SAM strictly outperforms SGD in terms of generalization error. Specifically, we show that SAM effectively mitigates noise learning in the early stages of training, enabling neural networks to learn features more efficiently.

Notation. We use lower case letters, lower case bold face letters, and upper case bold face letters to denote scalars, vectors, and matrices respectively. For a vector v = ( v 1 , · · · , v d ) ⊤ , we denote

by ∥ v ∥ 2 := ( ∑ d j =1 v 2 j ) 1 / 2 its 2 -norm. For two sequence { a k } and { b k } , we denote a k = O ( b k ) if | a k | ≤ C | b k | for some absolute constant C , denote a k = Ω( b k ) if b k = O ( a k ) , and denote a k = Θ( b k ) if a k = O ( b k ) and a k = Ω( b k ) . We also denote a k = o ( b k ) if lim | a k /b k | = 0 . Finally, we use ˜ O ( · ) 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 a set S is denoted by | S | .

2 Preliminaries

2.1 Data distribution

Our focus is on binary classification with label y ∈ {± 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 µ ∈ R d be a fixed vector representing the signal contained in each data point. Each data point ( x , y ) with input x = [ x (1) ⊤ , x (2) ⊤ , . . . , x ( P ) ⊤ ] ⊤ ∈ R P × d , x (1) , x (2) , . . . , x ( P ) ∈ R d and label y ∈ {-1 , 1 } is generated from a distribution D specified as follows:

The true label ̂ y is generated as a Rademacher random variable, i.e., P [ ̂ y = 1] = P [ ̂ y = -1] = 1 / 2 . The observed label y is then generated by flipping ̂ y with probability p where p < 1 / 2 , i.e., P [ y = y ] = 1 -p and P [ y = -y ] = p .

̂ ̂ 2. A noise vector ξ is generated from the Gaussian distribution N ( 0 , σ 2 p I ) , where σ 2 p is the variance.

One of x (1) , x (2) , . . . , x ( P ) is randomly selected and then assigned as y · µ , which represents the signal, while the others are given by ξ , which represents noises.

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

2.2 Neural Network and Training Loss

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

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

Here we consider ReLU activation function σ ( z ) = 1 ( z ≥ 0) z , w j,r ∈ R d denotes the weight for the r -th filter, and W j is the collection of model weights associated with F j for j = ± 1 . We use W to denote the collection of all model weights. Denote the training data set by S = { ( x i , y i ) } i ∈ [ n ] . We train the above CNN model by minimizing the empirical cross-entropy loss function

where ℓ ( z ) = log(1 + exp( -z )) .

2.3 Training Algorithm

Minibatch Stochastic Gradient Descent. For epoch t , the training data set S is randomly divided into H := n/B mini batches I t,b with batch size B ≥ 2 . The empirical loss for batch I t,b is defined

as L I t,b ( W ) = (1 /B ) ∑ i ∈I t,b ℓ ( y i f ( W , x i )) . Then the gradient descent update of the filters in the CNN can be written as

where the gradient of the empirical loss ∇ W L I t,b is the collection of ∇ w j,r L I t,b as follows

for all j ∈ {± 1 } and r ∈ [ m ] . Here we introduce a shorthand notation ℓ ′ ( t,b ) i = ℓ ′ [ y i · f ( W ( t,b ) , x i )] and assume the gradient of the ReLU activation function at 0 to be σ ′ (0) = 1 without loss of generality. We use ( t, b ) to denote epoch index t with mini-batch index b and use ( t ) as the shorthand of ( t, 0) . We initialize SGD by random Gaussian, where all entries of W (0) are sampled from i.i.d. Gaussian distributions N (0 , σ 2 0 ) , with σ 2 0 being the variance. From (3), we can infer that the loss landscape of the empirical loss is highly non-smooth because the ReLU function is not differentiable at zero. In particular, when ⟨ w ( t,b ) j,r , ξ ⟩ is close to zero, even a very small perturbation can greatly change the activation pattern σ ′ ( ⟨ w ( t,b ) j,r , ξ ⟩ ) and thus change the direction of ∇ w j,r L I t,b ( W ( t,b ) ) . This observation prevents the analysis technique based on the Taylor expansion with the Hessian matrix, and calls for a more sophisticated activation pattern analysis.

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

where the hyperparameter τ is called the perturbation radius. However, directly optimizing L SAM S ( W ) is computationally expensive. In practice, people use the following sharpness-aware minimization (SAM) algorithm (Foret et al., 2020; Zheng et al., 2021) to minimize L SAM S ( W ) efficiently,

When applied to SGD in (2), the gradient ∇ W L S in (5) is further replaced by stochastic gradient ∇ W L I t,b (Foret et al., 2020). The detailed algorithm description of SAM in shown in Algorithm 1.

Algorithm 1 Minibatch Sharpness Aware Minimization

Input: Training set S = ∪ n i =1 { ( x i , y i ) } , Batch size B , step size η > 0 , neighborhood size τ > 0 .

Initialize weights W (0) .

Randomly divide the training data set into H mini batches {I t,b } H -1

b =0 . for b = 0 , 1 , . . . , H -1 do

∇

∥∇

W

W

W

We calculate the perturbation

(

̂

ϵ

Update model parameters:

W

(

=

t,b t,b

)

+1)

τ

=

L

L

(

I

I

t,b

)

end for

Update model parameters: W ( t +1 , 0) = W ( t,H ) 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 σ 0 and learning rate η .

t,b t,b

(

W

(

-

W

(

(

t,b t,b

η

∇

)

W

)

)

)

∥

F

L

I

.

t,b

(

W

)

|

W

=

W

(

t,b

)

+

ϵ

(

̂

)

t,b

.

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

Dimension d is sufficiently large: d ≥ ˜ Ω ( max { nP -2 σ -2 p ∥ µ ∥ 2 2 , n 2 , P -2 σ -2 p Bm } ) .

Training sample size n and neural network width satisfy m,n ≥ ˜ Ω(1) .

The 2 -norm of the signal satisfies ∥ µ ∥ 2 ≥ ˜ Ω( Pσ p ) .

The noise rate p satisfies p ≤ 1 /C .

The standard deviation of Gaussian initialization σ 0 is appropriately chosen such that σ 0 ≤ ˜ O ( ( max { Pσ p d/ √ n, ∥ µ ∥ 2 }) -1 ) .

The learning rate η satisfies η ≤ ˜ O ( ( max { P 2 σ 2 p d 3 / 2 / ( Bm ) , P 2 σ 2 p d/B,n ∥ µ ∥ 2 / ( σ 0 B √ dm ) , nPσ p ∥ µ ∥ 2 / ( B 2 mϵ ) }) -1 ) .

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

Theorem 3.2 (Benign/harmful overfitting of SGD in training CNNs) . For any ϵ > 0 , under Condition 3.1, with probability at least 1 -δ there exists t = ˜ O ( η -1 ϵ -1 mnd -1 P -2 σ -2 p ) such that:

The training loss converges, i.e., L S ( W ( t ) ) ≤ ϵ .

When n ∥ µ ∥ 4 2 ≥ C 1 dP 4 σ 4 p , the test error L 0 -1 D ( W ( t ) ) ≤ p + ϵ .

When n ∥ µ ∥ 4 2 ≤ C 3 dP 4 σ 4 p , the test error L 0 -1 D ( W ( t ) ) ≥ p +0 . 1 .

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

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

3.1 Analysis of Mini-Batch SGD

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

Lemma 3.3. Let w ( t,b ) j,r for j ∈ {± 1 } , r ∈ [ m ] be the convolution filters of the CNN at the b -th batch of t -th epoch of gradient descent. Then there exist unique coefficients γ ( t,b ) j,r and ρ ( t,b ) j,r,i such that

Further denote ρ ( t,b ) j,r,i := ρ ( t,b ) j,r,i 1 ( ρ ( t,b ) j,r,i ≥ 0) , ρ ( t,b ) j,r,i := ρ ( t,b ) j,r,i 1 ( ρ ( t,b ) j,r,i ≤ 0) . Then

Note that (7) is a variant of (6): by decomposing the coefficient ρ ( t,b ) j,r,i into ρ ( t,b ) j,r,i and ρ ( t,b ) j,r,i , we can streamline our proof process. The normalization terms 1 P -1 , ∥ µ ∥ -2 2 , and ∥ ξ i ∥ -2 2 ensure that γ ( t,b ) j,r ≈ ⟨ w ( t,b ) j,r , µ ⟩ and ρ ( t,b ) j,r ≈ ( P -1) ⟨ w ( t,b ) j,r , ξ i ⟩ . Through signal-noise decomposition, we characterize the learning progress of signal µ using γ ( t,b ) j,r , and the learning progress of noise using ρ ( t,b ) j,r . This decomposition turns the analysis of SGD updates into the analysis of signal noise coefficients. Kou et al. (2023) extend this technique to the ReLU activation function as well as in the presence of label flipping noise. However, mini-batch SGD updates amplify the complications introduced by label flipping noise, making it more difficult to ensure learning. We have developed advanced methods for coefficient balancing and activation pattern analysis. These techniques will be thoroughly discussed in the sequel. The progress of signal learning is characterized by γ ( t,b ) j,r , whose update rule is as follows:

Here, I t,b represents the indices of samples in batch b of epoch t , S + denotes the set of clean samples where y i = ̂ y i , and S -represents the set of noisy samples where y i = -̂ y i . The updates of γ ( t,b ) j,r 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 ∑ i ∈I t,b ∩ S -ℓ ′ ( t,b ) i σ ′ ( ⟨ w ( t,b ) j,r , y i · µ ⟩ ) primarily drives model learning. Such unfavorable scenarios can be attributed to two possible factors. Firstly, the gradient of the loss ℓ ′ ( t,b ) i 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 t,b ∩ S -significantly outnumbers I t,b ∩ S + .

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

As ℓ ′ ( z 1 ) /ℓ ′ ( z 2 ) ≈ exp( z 2 -z 1 ) , we can upper bound ℓ ′ ( t,b 1 ) i /ℓ ′ ( t,b 2 ) k by y i · f ( W ( t,b 1 ) , x i ) -y k · f ( W ( t,b 2 ) , x k ) , which can be further upper bounded by ∑ r ρ ( t,b 1 ) y i ,r,i -∑ r ρ ( t,b 2 ) y i ,r,k with a small error. Therefore, the proof of (9) can be reduced to proving a uniform bound on the difference among ρ ( t,b ) y i ,r,i , i.e., ∑ m r =1 ρ ( t,b 1 ) y i ,r,i -∑ m r =1 ρ ( t,b 2 ) y k ,r,k ≤ κ, ∀ i, k .

However, achieving this uniform upper bound turns out to be challenging, since the updates of ρ ( t,b ) j,r,i 's are not evenly distributed across different batches within an epoch. Each mini-batch update utilizes only a portion of the samples, meaning that some ρ ( t,b ) y i ,r,i can increase or decrease much more than

the others. Therefore, the uniformly bounded difference can only be achieved after the entire epoch is processed. Consequently, we have to first bound the difference among ρ ( t,b ) y i ,r,i 's after each entire epoch, and then control the maximal difference within one epoch. The full batch (epoch) update rule is established as follows:

Here, b ( t ) i denotes the batch to which sample i belongs in epoch t , and ˜ S ( t,b ( t ) i ) i represents the parameters that learn ξ i at epoch t defined as

Therefore, the update of ∑ m r =1 [ ρ ( t, 0) y i ,r,i -ρ ( t, 0) y k ,r,k ] is indeed characterized by the activation pattern of parameters, which serves as the key technique for analyzing the full batch update of ∑ m r =1 [ ρ ( t, 0) y i ,r,i -ρ ( t, 0) y k ,r,k ] . However, analyzing the pattern of S ( t,b ) i directly is challenging since ⟨ w ( t,b ) y i ,r , ξ i ⟩ fluctuates in mini-batches without sample i . Therefore, we introduce the set series S ( t,b ) i as the activation pattern with certain threshold as follows:

The following lemma suggests that the set of activated parameters S ( t, 0) i is a non-decreasing sequence with regards to t , and the set of plain activated parameters ˜ S ( t,b ) i always include S ( t, 0) i . Consequently, S (0 , 0) i is always included in ˜ S ( t,b ) i , guaranteeing that ξ i can always be learned by some parameter. And this further makes sure the difference among ρ ( t,b ) y i ,r,i is bounded, as well as ℓ ′ ( t,b 1 ) i /ℓ ′ ( t,b 2 ) k ≤ C 2 . In the proof for SGD, we consider the learning period 0 ≤ t ≤ T ∗ , where T ∗ = η -1 poly( ϵ -1 , d, n, m ) is the maximum number of admissible iterations.

Lemma 3.4 (Informal Statement of Lemma C.8) . For all t ∈ [0 , T ∗ ] and b ∈ [ H ] , we have

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

Lemma 3.5 (Infomal Statement of Lemma B.6) . With high probability, for all T ∈ [ ˜ O (1) , T ∗ ] , there exist at least c 1 · T epochs among [0 , T ] , such that at least c 2 · H batches in each of these epochs satisfying the following condition:

After the first stage of T = Θ( η -1 m ( P -1) -2 σ -2 p d -1 ) epochs, we have γ ( T, 0) j,r = Ω ( n ∥ µ ∥ 2 2 ( P -1) 2 σ 2 p d ) .

The scale of γ ( T, 0) j,r guarantees that ⟨ w ( t,b ) j,r , µ ⟩ remains resistant to intra-epoch fluctuations. Consequently, this implies the sign of ⟨ w ( t,b ) j,r , µ ⟩ will persist unchanged throughout the entire epoch. Without loss of generality, we suppose that ⟨ w ( t,b ) j,r , µ ⟩ > 0 , then the update of γ ( t,b ) j,r can be written as follows:

As we have proved the balancing of logits ℓ ′ ( t,b ) i across batches, the progress analysis of γ ( t +1 , 0) j,r is established to characterize the signal learning of SGD.

4 Result for SAM

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

Theorem 4.1. For any ϵ > 0 , under Condition 3.1 with σ 0 = ˜ Θ( P -1 σ -1 p d -1 / 2 ) , choose τ = Θ ( m √ B Pσ p √ d ) . With probability at least 1 -δ , neural networks first trained with SAM with O ( η -1 ϵ -1 n -1 mB ∥ µ ∥ -2 2 ) iterations, then trained with SGD with ˜ O ( η -1 ϵ -1 mnd -1 P -2 σ -2 p ) iterations can find W ( t ) such that,

The training loss satisfies L S ( W ( t ) ) ≤ ϵ .

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

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, we have the following update rule of decomposition coefficients γ ( t,b ) j,r , ρ ( t,b ) j,r,i , ρ ( t,b ) j,r,i .

Lemma 4.2. The coefficients γ ( t,b ) j,r , ρ ( t,b ) j,r,i , ρ ( t,b ) j,r,i defined in Lemma 3.3 satisfy the following iterative equations for all r ∈ [ m ] , j ∈ {± 1 } and i ∈ [ n ] :

where I t,b denotes the sample index set of the b -th batch in the t -th epoch.

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

Lemma 4.3 (Informal Statement of Lemma D.5) . Suppose the Condition 3.1 holds with parameter choices in Theorem 4.1, if ⟨ w ( t,b ) j,r , ξ k ⟩ ≥ 0 , k ∈ I t,b and j = y k , then ⟨ w ( t,b ) j,r + ̂ ϵ ( t,b ) j,r , ξ k ⟩ < 0 .

By leveraging this intriguing property, we can derive a constant upper bound for the noise coefficients ρ ( t,b ) j,r,i by considering the following cases:

If ξ i is not in the current batch, then ρ ( t,b ) j,r,i will not be updated in the current iteration.

(a) Original SGD Test Error Heatmap

(b) SAM Test Error Heatmap

If ξ i is in the current batch, we discuss two cases:

(a) If ⟨ w ( t,b ) j,r , ξ i ⟩ ≥ 0 , then by Lemma 4.3, one can know that σ ′ ( ⟨ w ( t,b ) j,r + ̂ ϵ ( t,b ) j,r , ξ i ⟩ ) = 0 and thus ρ ( t,b ) j,r,i will not be updated in the current iteration.

(b) If ⟨ w ( t,b ) j,r , ξ i ⟩ ≤ 0 , then given that ⟨ w ( t,b ) j,r , ξ i ⟩ ≈ ρ ( t,b ) j,r,i and ρ ( t,b +1) j,r,i ≤ ρ ( t,b ) j,r,i + η ( P -1) 2 ∥ ξ i ∥ 2 2 Bm , we can assert that, provided η is sufficiently small, the term ρ ( t,b ) j,r,i can be upper bounded by a small constant.

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

5 Experiments

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

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

6 Related Work

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

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

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

7 Conclusion

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

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 arXiv: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 .

ALLEN-ZHU, Z. and LI, Y. (2022). Feature purification: How adversarial training performs robust deep learning. In 2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS) . IEEE.

ANDRIUSHCHENKO, M. and FLAMMARION, N. (2022). Towards understanding sharpness-aware minimization. In International Conference on Machine Learning . PMLR.

ANDRIUSHCHENKO, M., VARRE, A. V., PILLAUD-VIVIEN, L. and FLAMMARION, N. (2023). Sgd with large step sizes learns sparse features. In International Conference on Machine Learning . PMLR.

BAHRI, D., MOBAHI, H. and TAY, Y. (2021). Sharpness-aware minimization improves language model generalization. arXiv preprint arXiv:2110.08529 .

BARTLETT, P. L., LONG, P. M. and BOUSQUET, O. (2022). The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. arXiv preprint arXiv:2210.01513 .

BARTLETT, P. L., LONG, P. 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-aware minimization: An implicit regularization perspective. arXiv preprint arXiv:2302.11836 .

BEHDIN, K., SONG, Q., GUPTA, A., DURFEE, D., ACHARYA, A., KEERTHI, S. and MAZUMDER, R. (2022). Improved deep neural network generalization using m-sharpness-aware minimization. arXiv preprint arXiv:2212.04343 .

CAO, 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 arXiv:2209.09315 .

CHEN, X., HSIEH, C.-J. and GONG, B. (2021). When vision transformers outperform resnets without pre-training or strong data augmentations. arXiv preprint arXiv:2106.01548 .

DEVROYE, L., MEHRABIAN, A. and REDDAD, T. (2018). The total variation distance between high-dimensional gaussians. arXiv preprint arXiv:1810.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. In 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 arXiv:2303.04145 .

KWON, J., KIM, J., PARK, H. and CHOI, I. K. (2021). Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning . PMLR.

LI, Y., WEI, C. and MA, T. (2019). Towards explaining the regularization effect of initial large learning rate in training neural networks. In Advances in Neural Information Processing Systems .

LI, Z., WANG, T. and ARORA, S. (2021a). What happens after sgd reaches zero loss?-a mathematical framework. arXiv preprint arXiv:2110.06914 .

LI, Z., ZHOU, Z.-H. and GRETTON, A. (2021b). Towards an understanding of benign overfitting in neural networks. arXiv preprint arXiv:2106.03212 .

LIU, Y., MAI, S., CHEN, X., HSIEH, C.-J. and YOU, Y. (2022a). Towards efficient and scalable sharpness-aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) .

LIU, Y., MAI, S., CHEN, X., HSIEH, C.-J. and YOU, Y. (2022b). Towards efficient and scalable sharpness-aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition .

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 ZHONG, Y. (2022). The interpolation phase transition in neural networks: Memorization and generalization under lazy training. The Annals of Statistics 50 2816-2847.

OLSHAUSEN, B. A. and FIELD, D. J. (1997). Sparse coding with an overcomplete basis set: A strategy employed by v1? Vision research 37 3311-3325.

SHEN, R., BUBECK, S. and GUNASEKAR, S. (2022). Data augmentation as feature manipulation: a story of desert cows and grass cows. arXiv preprint arXiv:2203.01572 .

VERSHYNIN, R. (2018). High-Dimensional 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, R. and MAO, Y. (2021). Regularizing neural networks via adversarial model perturbation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition .

ZHUANG, J., GONG, B., YUAN, L., CUI, Y., ADAM, H., DVORNEK, N., TATIKONDA, S., DUNCAN, J. and LIU, T. (2022). Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.08065 .

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 in the early stage of training. Is SAM also effective if we add the algorithm in the middle of the training process? We conduct experiments on the ImageNet dataset with ResNet50. We choose the batch size as 1024 , and the model is trained for 90 epochs with the best learning rate in grid search { 0 . 01 , 0 . 03 , 0 . 1 , 0 . 3 } . The learning rate schedule is 10 k steps linear warmup, then cosine decay. As shown in Table 1, the earlier SAM is introduced, the more pronounced its effectiveness becomes.

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

A.2 Discussion on the Stochastic Gradient Descent

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

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

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

0.0

5.0

Strength of the Signal

1000

11000

Dimension

0.0

0.2

0.4

0.6

0.8

1.0

0.0

5.0

Strength of the Signal

1000

11000

Dimension

0.0

0.2

0.4

0.6

0.8

1.0

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

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

B Preliminary Lemmas

Lemma B.1 (Lemma B.4 in Kou et al. (2023)) . Suppose that δ > 0 and d = Ω(log(6 n/δ )) . Then with probability at least 1 -δ ,

for all i, i ′ ∈ [ n ] .

Lemma B.2 (Lemma B.5 in Kou et al. (2023)) . Suppose that d = Ω(log( mn/δ )) , m = Ω(log(1 /δ )) . Then with probability at least 1 -δ ,

for all r ∈ [ m ] , j ∈ {± 1 } and i ∈ [ n ] . Moreover,

for all j ∈ {± 1 } and i ∈ [ n ] .

Lemma B.3. Let S ( t,b ) i denote { r : ⟨ w ( t,b ) y i ,r , ξ i ⟩ > σ 0 σ p √ d/ √ 2 ⟩} . Suppose that δ > 0 and m ≥ 50 log(2 n/δ ) . Then with probability at least 1 -δ ,

0.0

5.0

Strength of the Signal

1000

11000

Dimension

0.0

0.2

0.4

0.6

0.8

1.0

0.0

5.0

Strength of the Signal

1000

11000

Dimension

0.0

0.2

0.4

0.6

0.8

1.0

0.0

5.0

Strength of the Signal

1000

11000

Dimension

0.0

0.2

0.4

0.6

0.8

1.0

0.0

5.0

Strength of the Signal

1000

11000

Dimension

0.0

0.2

0.4

0.6

0.8

1.0

Proof of Lemma B.3. Since ⟨ w (0 , 0) y i ,r , ξ i ⟩ ∼ N (0 , σ 2 0 ∥ ξ i ∥ 2 2 ) , we have

where Φ( · ) is CDF of the standard normal distribution. Note that | S (0 , 0) i | = ∑ m r =1 1 [ ⟨ w (0 , 0) y i ,r , ξ i ⟩ > σ 0 σ p √ d/ √ 2] and P ( ⟨ w (0 , 0) y i ,r , ξ i ⟩ > σ 0 σ p √ d/ √ 2 ) ≥ Φ( -1) , then by Hoeffding's inequality, with probability at least 1 -δ/n , we have

Therefore, as long as 0 . 2 √ m Φ( -1) ≥ √ log(2 n/δ ) 2 , by applying union bound, with probability at least 1 -δ , we have

Lemma B.4. Let S ( t,b ) j,r denote { i ∈ [ n ] : y i = j, ⟨ w ( t,b ) y i ,r , ξ i ⟩ > σ 0 σ p √ d/ √ 2 } . Suppose that δ > 0 and n ≥ 32 log(4 m/δ ) . Then with probability at least 1 -δ ,

Proof of Lemma B.4. Since ⟨ w (0 , 0) j,r , ξ i ⟩ ∼ N (0 , σ 2 0 ∥ ξ i ∥ 2 2 ) , we have

where Φ( · ) is CDF of the standard normal distribution.

Note that | S (0 , 0) j,r | = ∑ n i =1 1 [ y i = j ] 1 [ ⟨ w (0) j,r , ξ i ⟩ > σ 0 σ p √ d/ √ 2] and P ( y i = j, ⟨ w (0) j,r , ξ i ⟩ > σ 0 σ p √ d/ √ 2) ≥ Φ( -1) / 2 , then by Hoeffding's inequality, with probability at least 1 -δ/ 2 m , we have

Therefore, as long as Φ( -1) / 4 ≥ √ log(4 m/δ ) 2 n , by applying union bound, we have with probability at least 1 -δ ,

Lemma B.5 (Lemma B.3 in Kou et al. (2023)) . For | S + ∩ S y | and | S -∩ S y | where y ∈ {± 1 } , it holds with probability at least 1 -δ ( δ > 0) that

Lemma B.6. It holds with probability at least 1 -δ , for all T ∈ [ log(2 T ∗ /δ ) c 2 3 , T ∗ ] and y ∈ {± 1 } , there exist at least c 3 · T epochs among [0 , T ] , such that at least c 4 · H batches in these epochs, satisfy

Proof. Let

E 1 ,t := { In epoch t , there are at least c 2 · n B batches such that 16 holds for y = 1 } ,

E 1 ,t,b := { In epoch t natch b , 16 holds for y = 1 } .

First let n big enough, then we have S + ∩ S y ∈ [ 3(1 -p ) n 8 , 5(1 -p ) n 8 ] . We consider the first c 1 H batches. At the time we are starting to sample h -th batch in the first c 1 H batches, suppose there are n 1 samples that belong to S + ∩ S y , and there are n 2 samples that don't belong to S + ∩ S y . Then n 1 ≥ 3(1 -p ) n 8 -c 1 n ≥ 5(1 -p ) n 16 and n 2 ≥ 3(1 -p ) n 8 -c 1 n ≥ 5(1 -p ) n 16 .

Then, the probability that there are less than c 1 c 2 H batches in first c 1 H batches such that 16 holds is:

Choose H 0 such that c 1 c 2 H 0 · (2 c 2 ) c 1 c 2 H 0 (1 -2 c 2 ) c 1 H 0 -c 1 c 2 H 0 = 1 -2 c 3 , then as long as H ≥ H 0 , with probability c 3 , there are at least c 1 c 2 H batches in first c 1 H batches such that 16 holds. Then P [ E 1 ,t ] ≥ 2 c 3 . Therefore,

Let T ≥ log(2 T ∗ /δ ) 2 c 2 3 . Then, with probability at least 1 -δ/ (2 T ∗ ) ,

Let c 4 = c 1 c 2 . Thus, there are at least c 3 T ∗ epochs, such that they have at least c 4 H batches satisfying EquationB.6. This also holds for y = -1 . Taking a union bound to get the result.

C Proofs for SGD

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

C.1 Signal-noise Decomposition Coefficient Analysis

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

Let I t,b denote the set of indices of randomly chosen samples at epoch t batch b , and |I t,b | = B , then the update rule is:

C.1.1 Iterative Expression for Decomposition Coefficient Analysis

Lemma C.1. The coefficients γ ( t,b ) j,r , ρ ( t,b ) j,r,i , ρ ( t,b ) j,r,i defined in Lemma 3.3 satisfy the following iterative equations:

for all r ∈ [ m ] , j ∈ {± 1 } and i ∈ [ n ] .

Proof. First, we iterate the gradient descent update rule t epochs plus b batches and get

According to the definition of γ ( t ) j,r and ρ ( t ) j,r,i ,

Since ξ i and µ are linearly independent with probability 1 , we have the unique representation as follows:

Since we define ρ ( t,b ) j,r,i := ρ ( t,b ) j,r,i 1 ( ρ ( t,b ) j,r,i ≥ 0) , ρ ( t ) j,r,i := ρ ( t,b ) j,r,i 1 ( ρ ( t,b ) j,r,i ≤ 0) , we obtain

And the iterative update equations (19), (20), and (21) follow directly.

C.1.2 Scale of Decomposition Coefficients

We first define T ∗ = η -1 poly( ϵ -1 , d, n, m ) and

By Lemma B.2 and Condition 3.1, β can be bounded as

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

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

Proposition C.2. Under Assumption 3.1, for (0 , 0) ≤ ( t, b ) ≤ ( T ∗ , 0) , we have that

and there exists a positive constant C ′ such that

for all r ∈ [ m ] , j ∈ {± 1 } and i ∈ [ n ] , where ̂ γ := n · SNR 2 .

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

Lemma C.3. Under Assumption 3.1, suppose (27) , (28) and (29) hold after b -th batch of t -th epoch. Then, for all r ∈ [ m ] , j ∈ {± 1 } and i ∈ [ n ] ,

̸

Proof of Lemma C.3. First, for any time ( t, b ) ≥ (0 , 0) , we have from the following decomposition by definitions,

According to Lemma B.1, we have

where the first inequality is by triangle inequality, the second inequality is by Lemma B.1, the equality is by SNR = ∥ µ ∥ 2 / (( P -1) σ p √ d ) , and the last inequality is by (27), (28). It follows that

̸

Then, for j = y i and any t ≥ 0 , we have

̸

̸

where the second equality is due to ρ ( t,b ) j,r,i = 0 for j = y i . Next, we have

̸

̸

̸

where the first inequality is by triangle inequality; the second inequality is by Lemma B.1; the equality is by SNR = ∥ µ ∥ 2 /σ p √ d ; the third inequality is by (28) and (29); the forth inequality is by SNR ≤ 1 / √ 8 C ′ 2 . Therefore, for j = y i , we have

Similarly, we have for y i = j that

̸

̸

and

̸

̸

̸

where the second inequality is by (27) and (29), and the third inequality is by SNR ≤ 1 / √ 8 C ′ 2 . Therefore, for j = y i , we have

̸

Lemma C.4. Under Condition 3.1, suppose (27) , (28) and (29) hold after b -th batch of t -th epoch. Then, for all j = y i , j ∈ {± 1 } and i ∈ [ n ] , F j ( W ( t,b ) j , x i ) ≤ 0 . 5 .

Proof of Lemma C.4. According to Lemma C.3, we have

where the second inequality is by (30), (31) and (32); the third inequality is due to the definition of β and ρ ( t,b ) j,r,i < 0 ; the third inequality is by (25) and -γ ( t,b ) j,r ≤ 1 12 .

Lemma C.5. Under Condition 3.1, suppose (27) , (28) and (29) hold at b -th batch of t -th epoch. Then, it holds that

Proof of Lemma C.5. According to (32) in Lemma C.3, we have

where the second inequality is due to ρ ( t,b ) y i ,r,i ≥ 0 , the third inequality is due to β < 1 / 8 and 5 n √ log(6 n 2 /δ ) /d · α < 1 / 8 by Condition 3.1.

For the second equation, the first inequality holds naturally since z ≤ σ ( z ) . For the inequality, if ⟨ w ( t ) y i ,r , ξ i ⟩ ≤ 0 , we have

And if ⟨ w ( t,b ) y i ,r , ξ i ⟩ > 0 , we have

Lemma C.6 (Lemma C.6 in Kou et al. (2023)) . Let g ( z ) = ℓ ′ ( z ) = -1 / (1 + exp( z )) , then for all z 2 -c ≥ z 1 ≥ -1 where c ≥ 0 we have that

Lemma C.7. For any iteration t ∈ [0 , T ∗ ) and b, b 1 , b 2 ∈ [ H ] , we have the following statements hold:

Proof. For the first statement,

where the first inequality follows from the iterative update rule of ρ ( t,b ) j,r,i , the second inequality is due to Lemma B.2, and the last inequality is due to Condition 3.1.

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

Therefore,

where the first inequality is due to Lemma B.1, and the second inequality is due to Condition 3.1. For the third statement. Let r ∗ ∈ S ( t, 0) , then

i

where the first inequality is due to the second statement, and the second inequality is due to the definition of ˜ S t, 0 i . Therefore, r ∗ ∈ ˜ S ( t,b ) i and S ( t,b ) i ⊆ ˜ S ( t,b ) i . The fourth statement can be obtained similarly.

Lemma C.8. Under Assumption 3.1, suppose (27) , (28) and (29) hold for any iteration ( t ′ , b ′ ) ≤ ( t, 0) . Then, the following conditions also hold for ∀ t ′ ≤ t and ∀ b ′ , b ′ 1 , b ′ 2 ∈ [ H ] :

Here we take κ and C 1 as 10 and 5 respectively.

Proof of Lemma C.8. We prove Lemma C.8 by induction. When t ′ = 0 , the fourth and fifth conditions hold naturally by Lemma B.3 and B.4.

For the first condition, since we have ρ (0 , 0) j,r,i = 0 for any j, r, i according to (26), it is straightforward that ∑ m r =1 [ ρ (0 , 0) y i ,r,i -ρ (0 , 0) y k ,r,k ] = 0 for all i, k ∈ [ n ] . So the first condition holds for t ′ = 0 .

For the second condition, we have

where the first inequality is by F j ( W (0 , 0) j , x i ) > 0 , the second inequality is due to (23), and the third inequality is due to (25).

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

Now suppose there exists ( ˜ t, ˜ b ) ≤ ( t, b ) such that these five conditions hold for any (0 , 0) ≤ ( t ′ , b ′ ) < ( ˜ t, ˜ b ) . We aim to prove that these conditions also hold for ( t ′ , b ′ ) = ( ˜ t, ˜ b ) .

We first show that, for any 0 ≤ t ′ ≤ t and 0 ≤ b ′ 1 , b ′ 2 ≤ b , y i · f ( W ( t ′ ,b ′ 1 ) , x i ) -y k · f ( W ( t ′ ,b ′ 2 ) , x k ) can be approximated by 1 m ∑ m r =1 [ ρ ( t ′ ,b ′ 1 ) y i ,r,i -ρ ( t ′ ,b ′ 2 ) y k ,r,k ] with a small constant approximation error. We begin by writing out

where all the equalities are due to the network definition. Then we bound I 1 , I 2 and I 3 . For | I 1 | , we have the following upper bound by Lemma C.4:

For | I 2 | , we have the following upper bound:

where the second inequality is due to (30), the second inequality is due to the definition of β and (29), the third inequality is due to Condition 3.1 and (25).

For I 3 , we have the following upper bound

where the first inequality is due to Lemma C.5, the second inequality is due to Lemma C.3, the third inequality is due to 5 √ log(6 n 2 /δ ) /dnα ≤ 1 / 8 according to Condition 3.1.

Similarly, we have the following lower bound

where the first inequality is due to Lemma C.5, the second inequality is due to Lemma C.3, the third inequality is due to 5 √ log(6 n 2 /δ ) /dnα ≤ 1 / 8 according to Condition 3.1.

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

which is equivalent to

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

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

for all j ∈ {± 1 } , r ∈ [ m ] , i ∈ [ n ] , (0 , 0) ≤ ( t, b ) < [ T ∗ , 0] . It follows that

for all i, k ∈ [ n ] and 0 ≤ t ≤ T ∗ , b < H .

If ˜ b ∈ { 1 , 2 , · · · , H -1 } , then the first statement for ( t ′ , b ′ ) = ( ˜ t, ˜ b ) and for the last ( t ′ , b ′ ) < ( ˜ t, ˜ b ) are the same. Otherwise, if ˜ b = 0 , we consider two separate cases: ∑ m r =1 [ ρ ( ˜ t -1 , 0) y i ,r,i -ρ ( ˜ t -1 , 0) y k ,r,k ] ≤ 0 . 9 κ and ∑ m r =1 [ ρ ( ˜ t -1 , 0) y i ,r,i -ρ ( ˜ t -1 , 0) y k ,r,k ] > 0 . 9 κ .

When ∑ m r =1 [ ρ ( ˜ t -1 , 0) y i ,r,i -ρ ( ˜ t -1 , 0) y k ,r,k ] ≤ 0 . 9 κ , we have

where the first inequality is due to ℓ ′ ( ˜ t -1 ,b ( ˜ t -1) i ) i < 0 ; the second inequality is due to ∣ ∣ S ( ˜ t -1 ,b ( ˜ t -1) i ) i ∣ ∣ ≤ m and -ℓ ′ ( ˜ t -1 ,b ( ˜ t -1) i ) i < 1 ; the third inequality is due to Condition 3.1.

On the other hand, for when ∑ m r =1 [ ρ ( ˜ t -1 , 0) y i ,r,i -ρ ( ˜ t -1 , 0) y k ,r,k ] > 0 . 9 κ , we have from the (38) that

where the second inequality is due to κ = 10 . Thus, according to Lemma C.6, we have

Since S ( ˜ t -1 , 0) i ⊆ ˜ S ( ˜ t -1 ,b ( ˜ t -1) i ) i , we have ∣ ∣ ∣ ∣ ˜ S ( ˜ t -1 ,b ( ˜ t -1) k ) k ∣ ∣ ∣ ∣ ≥ 0 . 8Φ( -1) m according to the fourth condition. Also we have that | S ( ˜ t -1 ,b ( ˜ t -1) i ) i | ≤ m . It follows that

According to Lemma B.1, under event E prelim , we have

Note that d = Ω(log(6 n/δ )) from Condition 3.1, it follows that

Then we have

which completes the proof of the first hypothesis at iteration ( t ′ , b ′ ) = ( ˜ t, ˜ b ) . Next, by applying the approximation in (38), we are ready to verify the second hypothesis at iteration ( ˜ t, ˜ b ) . In fact, for any ( t ′ , b ′ 1 ) , ( t ′ , b ′ 2 ) ≤ ( ˜ t, ˜ b ) , we have

where the first inequality is by (38); the last inequality is by induction hypothesis and taking κ as 10 and C 1 as 5.

And the third hypothesis directly follows by noting that, for any ( t ′ , b ′ 1 ) , ( t ′ , b ′ 2 ) ≤ ( ˜ t, ˜ b ) ,

For the fourth hypothesis, If ˜ b ∈ { 1 , 2 , · · · , H -1 } , then the first statement for ( t ′ , b ′ ) = ( ˜ t, ˜ b ) and for the last ( t ′ , b ′ ) < ( ˜ t, ˜ b ) are the same. Otherwise, if ˜ b = 0 , according to the gradient descent rule, we have

̸

̸

for any r ∈ S ( ˜ t -1 , 0) i , where the last equality is by ⟨ w ( ˜ t -1 ,b ( ˜ t -1) i ) y i ,r , ξ i ⟩ > 0 . Then we respectively estimate I 4 , I 5 , I 6 . For I 4 , according to Lemma B.1, we have

For I 5 , we have following upper bound

̸

where the first inequality is due to triangle inequality, the second inequality is due to σ ′ ( z ) ∈ { 0 , 1 } , the third inequality is due to Lemma B.1, the forth inequality is due to the third hypothesis at epoch ˜ t -1 .

For I 6 , we have following upper bound

where the first inequality is by triangle inequality; the second inequality is due to σ ′ ( z ) ∈ { 0 , 1 } ; the third inequality is by Lemma B.1; the last inequality is due to the third hypothesis at epoch ˜ t -1 .

Since d ≥ max { 32 C 2 2 n 2 · log(6 n 2 /δ ) , 4 C 2 n ∥ µ ∥ σ -1 p √ 2 log(6 n/δ ) } , we have -( P -1)I 4 max { ( P -1) | I 5 | / 2 , | I 6 | / 2 } and hence -( P -1)I 4 ≥ ( P -1) | I 5 | + | I 6 | . It follows that

for any r ∈ S ( ˜ t -1 , 0) i . Therefore, S (0 , 0) i ⊆ S ( ˜ t -1 , 0) i ⊆ S ( ˜ t, 0) i . And it directly follows by Lemma B.3 that | S ( ˜ t, 0) i | ≥ 0 . 8 m Φ( -1) , ∀ i ∈ [ n ] .

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

̸

̸

for any i ∈ S ( ˜ t -1 , 0) j,r , where the equality holds due to ⟨ w ( ˜ t -1 ,b ( ˜ t -1) i ) j,r , ξ i ⟩ > 0 and y i = j . By applying the same technique used in the proof of the fourth hypothesis, it follows that

for any i ∈ S ( ˜ t -1 , 0) j,r . Thus, we have S (0 , 0) j,r ⊆ S ( ˜ t -1 , 0) j,r ⊆ S ( ˜ t, 0) j,r . And it directly follows by Lemma B.4 that | S ( ˜ t, 0) j,r | ≥ n Φ( -1) / 4 .

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

̸

Firstly, We prove that (28) exists at iteration ( ˜ t, ˜ b ) , i.e., ρ ( ˜ t, ˜ b ) j,r,i ≥ -β -10 √ log(6 n 2 /δ ) /d · nα for any r ∈ [ m ] , j ∈ {± 1 } and i ∈ [ n ] . Notice that ρ ( ˜ t, ˜ b ) j,r,i = 0 for j = y i , therefore we only need to consider the case that j = y i . We also only need to consider the case of ˜ b = b ( ˜ t ) i +1 since ρ ( ˜ t, ˜ b ) j,r,i doesn't change in other cases according to (21).

When ρ ( ˜ t,b ( ˜ t ) i ) j,r,t < -0 . 5 β -5 √ log(6 n 2 /δ ) /d · nα , by (32) in Lemma C.3 we have that

and thus

where the last inequality is by induction hypothesis.

where the first inequality is by ℓ ′ ( t,b ( ˜ t ) i ) i ∈ ( -1 , 0) and ∥ ξ i ∥ 2 2 ≤ (3 / 2) σ 2 p d by Lemma B.1; the second inequality is due to 5 √ log(6 n 2 /δ ) /d · nα ≥ 3 ησ 2 p d/ (2 Bm ) by Condition 3.1.

Next we prove (27) holds for ( ˜ t, ˜ b ) . We only need to consider the case of j = y i . Consider

̸

where the last inequality is by F j ( W ( ˜ t, ˜ b ) j , x i ) ≤ 0 . 5 for j = y i according to Lemma C.4. Now recall the iterative update rule of ρ ( t,b ) j,r,i :

Let ( t j,r,i , b j,r,i ) be the last time before ( ˜ t, ˜ b ) that ρ ( t j,r,i ,b j,r,i ) j,r,i ≤ 0 . 5 α . Then by iterating the update rule from ( t j,r,i , b j,r,i ) to ( ˜ t, ˜ b ) , we get

We first bound I 7 as follows:

where the first inequality is by ℓ ′ ( t j,r,i ,b j,r,i ) i ∈ ( -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 α = 4log( T ∗ ) and T ∗ ≥ e .

Second, we bound I 8 . For ( t j,r,i , b j,r,i ) < ( t, b ) < ( ˜ t, ˜ b ) and y i = j , we can lower bound the inner product ⟨ w ( t,b ) j,r , ξ i ⟩ as follows

where the first inequality is by (31) in Lemma C.3; the second inequality is by ρ ( t,b ) j,r,i > 0 . 5 α and ⟨ w (0 , 0) j,r , ξ i ⟩ ≥ -0 . 5 β/ ( P -1) due to the definition of t j,r,i and β ; the last inequality is by β ≤ 1 / 8 ≤ 0 . 1 α and 5 √ log(6 n 2 /δ ) /d · nα ≤ 0 . 2 α by Condition 3.1.

Thus, plugging the lower bounds of ⟨ w ( t,b ) j,r , ξ i ⟩ into I 8 gives

where the first inequality is by (40); the second inequality is by (42); the third inequality is by α = 4log( T ∗ ) ; the fourth inequality is by Condition 3.1; the last inequality is by log( T ∗ ) ≥ 1 and α = 4log( T ∗ ) . Plugging the bound of I 7 , I 8 into (41) completes the proof for ρ .

For the upper bound of (29), we prove a augmented hypothesis that there exists a i ∗ ∈ [ n ] with y i ∗ = j such that for 1 ≤ t ≤ T ∗ we have that γ ( t, 0) j,r /ρ j,r,i ∗ ≤ C ′ ̂ γ . Recall the iterative update rule of γ ( t,b ) j,r and ρ ( t,b ) j,r,i , we have

According to the fifth statement of Lemma C.8, for any i ∗ ∈ S (0 , 0) j,r it holds that j = y i ∗ and ⟨ w ( t,b ) j,r , ξ i ∗ ⟩ ≥ 0 for any ( t, b ) ≤ ( ˜ t, ˜ b ) . Thus, we have

For the update rule of γ ( t,b ) j,r , according to Lemma C.8, we have

Then, we have

where the last inequality is by γ ( ˜ t -1 , 0) j,r /ρ ( ˜ t -1 , 0) j,r,i ∗ ≤ 2 C 2 ̂ γ = 2 C 2 n ∥ µ ∥ 2 2 / ( P -1) 2 σ 2 p d . Therefore,

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

Thus, by taking C ′ = 4 C 2 , we have γ ( ˜ t,b ) j,r /ρ ( ˜ t,b ) j,r,i ∗ ≤ C ′ ̂ γ .

On the other hand, when ( t, b ) < ( log(2 T ∗ /δ ) 2 c 2 3 , 0) , we have

where the first inequality is due to update rule of γ t,b j,r , and the second inequality is due to Condition 3.1.

When ( t, b ) ≥ ( log(2 T ∗ /δ ) 2 c 2 3 , 0) , According to Lemma B.6, we have γ ( t,b ) j,r ≥ ∑ ( t ′ ,b ′ ) < ( t,b ) η Bm [ min i,b ′ ℓ ′ ( t ′ ,b ′ ) i min {|I t ′ ,b ′ ∩ S + ∩ S -1 | , |I t ′ ,b ′ ∩ S + ∩ S 1 |} -max i,b ′ ℓ ′ ( t ′ ,b ′ ) i |I t ′ ,b ′ ∩ S -| ] · ∥ µ ∥ 2 2 ≥ η Bm ( t -1 ∑ t ′ =0 ( c 3 c 4 H B 4 min i,b ′ ℓ ′ ( t ′ ,b ′ ) i -nq max i,b ′ ℓ ′ ( t ′ ,b ′ ) i ) -nq max i,b ′ ℓ ′ ( t,b ′ ) i )) ∥ µ ∥ 2 2 ≥ 0 , where the first inequality is due to the update rule of γ ( t,b ) j,r , the second inequality is due to Lemma B.6, and the third inequality is due to Condition 3.1.

C.2 Decoupling with a Two-stage Analysis

C.2.1 First Stage

Lemma C.9. There exist

where C 3 = Θ(1) is a large constant and C 4 = Θ(1) is a small constant, such that

Proof of Lemma C.9. By Proposition C.2, we have that ρ ( t,b ) j,r,i ≥ -β -10 n √ log(6 n 2 /δ ) d α for all j ∈ {± 1 } , r ∈ [ m ] , i ∈ [ n ] and (0 , 0) ≤ ( t, b ) ≤ ( T ∗ , 0) . According to Lemma B.2, for β we have

where the last equality is by the first condition of Condition 3.1. Since ρ ( t,b ) j,r,i ≤ 0 , we have that

Next, for the growth of γ ( t ) j,r , we have following upper bound

where the inequality is by | ℓ ′ | ≤ 1 . Note that γ (0 , 0) j,r = 0 and recursively use the inequality tB + b times we have

Since n · SNR 2 = n ∥ µ ∥ 2 2 / ( ( P -1) 2 σ 2 p d ) = ̂ γ , we have

And it follows that

for all (0 , 0) ≤ ( t, b ) ≤ ( T 1 , 0) .

For ρ ( t ) j,r,i , recall from (20) that

According to Lemma C.8, for any r ∗ ∈ S (0 , 0) i = { r ∈ [ m ] : ⟨ w (0) y i ,r , ξ i ⟩ > σ 0 σ p √ d/ √ 2 } , we have ⟨ w ( t,b ) y i ,r ∗ , ξ i ⟩ > 0 for all (0 , 0) ≤ ( t, b ) ≤ ( T ∗ , 0) and hence

For each i , we denote by T ( i ) 1 the last time in the period [0 , T 1 ] satisfying that ρ ( t, 0) y i ,r ∗ ,i ≤ 2 . Then for (0 , 0) ≤ ( t, b ) < ( T ( i ) 1 , 0) , max j,r {| ρ ( t,b ) j,r,i | , | ρ ( t,b ) j,r,i |} = O (1) and max j,r γ ( t,b ) j,r = O (1) . Therefore, we know that F -1 ( W ( t,b ) , x i ) , F +1 ( W ( t,b ) , x i ) = O (1) . Thus there exists a positive constant C such that -ℓ ′ ( t,b ) i ≥ C ≥ C 2 for 0 ≤ t ≤ T ( i ) 1 .

Then we have

Therefore, ρ ( t, 0) y i ,r ∗ ,i will reach 2 within

iterations for any r ∗ ∈ S (0 , 0) i , where C 3 can be taken as 4 /C .

Next, we will discuss the lower bound of the growth of γ ( t,b ) j,r . For ρ ( t,b ) j,r,i , we have

According to (44) and ρ (0 , 0) j,r,i = 0 , it follows that

Therefore, max j,r,i ρ ( t,b ) j,r,i will be smaller than 1 and γ ( t,b ) j,r smaller than Θ( n ∥ µ ∥ 2 2 / ( P -1) 2 σ 2 p d ) = Θ( n · SNR 2 ) = Θ( ̂ γ ) = O (1) within

iterations, where C 4 can be taken as 2 / 3 . Therefore, we know that F -1 ( W ( t,b ) , x i ) , F +1 ( W ( t,b ) , x i ) = O (1) in (0 , 0) ≤ ( t, b ) ≤ ( T 2 , 0) . Thus, there exists a positive constant C such that -ℓ ′ ( t,b ) i ≥ C for 0 ≤ t ≤ T 2 .

Recall that we denote { i ∈ [ n ] | y i = y } as S y , and we have the update rule

For the growth of γ ( t,b ) j,r , if ⟨ w ( t,b ) j,r , µ ⟩ ≥ 0 , we have

Similarly, if ⟨ w ( t,b ) j,r , µ ⟩ < 0 ,

Therefore, for t ∈ [ T 2 , T 1 ] , we have

where the second inequality is due to Lemma B.6, the third inequality is due to q < C 4 Cc 3 c 4 8 C 3 in Condition 3.1.

And it follows directly from (45) that

which completes the proof.

C.2.2 Second Stage

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

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

max j,r,i | ρ ( T 1 , 0) j,r,i | = max { β, O ( n √ log( n/δ ) log( T ∗ ) / √ d ) } .

where ̂ γ = n · SNR 2 . Now we choose W ∗ as follows

Lemma C.10. Under the same conditions as Theorem 3.2, we have that ∥ W ( T 1 , 0) -W ∗ ∥ F ≤ ˜ O ( m 1 / 2 n 1 / 2 ( P -1) -1 σ -1 p d -1 / 2 (1 + max { β, n √ log( n/δ ) log( T ∗ ) / √ d } )) .

Proof.

where the first inequality is by triangle inequality, the second inequality and the first equality are by our decomposition of W ( T 1 , 0) , W ∗ and Lemma B.1; the second equality is by n · SNR 2 = Θ( ̂ γ ) and SNR = ∥ µ ∥ / ( P -1) σ p d 1 / 2 ; the third equality is by n 1 / 2 · SNR = O (1) .

Lemma C.11. Under the same conditions as Theorem 3.2, we have that

Proof of Lemma C.11. Recall that f ( W ( t,b ) ) = (1 /m ) ∑ j,r j · [ σ ( ⟨ w ( t,b ) j,r , ̂ y i µ ⟩ ) + ( P -1) σ ( ⟨ w ( t,b ) j,r , ξ i ⟩ )] , thus we have

̸

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

For I 10 , I 11 , I 12 , I 14 , we have that

For j = y i and r ∈ S (0) i , according to Lemma C.3, we have

where the first inequality is by Lemma C.3; the second inequality is by 5 n √ log(4 n 2 /δ ) d ≤ 0 . 5 ; and the last inequality is by β < 1 . 5 . Therefore, for I 9 , according to the fourth statement of Proposition C.8, we have

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

which completes the proof.

Lemma C.12. Under Assumption 3.1, for (0 , 0) ≤ ( t, b ) ≤ ( T ∗ , 0) , the following result holds.

Proof. We first prove that

Without loss of generality, we suppose that ̂ y i = 1 . Then we have that

where the first and second inequalities are by triangle inequality, the third inequality is by Lemma B.1.

Then we upper bound the gradient norm ∥∇ L S ( W ( t,b ) ) ∥ F as:

where the first inequality is by triangle inequality, the second inequality is by (50), the third inequality is by Cauchy-Schwartz inequality and the last inequality is due to the property of the cross entropy loss -ℓ ′ ≤ ℓ .

Lemma C.13. Under the same conditions as Theorem 3.2, we have that for all ( T , 0) ≤ ( t, b ) ≤ ( T , 0)

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 to Lemma C.12.

Lemma C.14. Under the same conditions as Theorem 3.2, with probability 1 -δ , we have that

Proof.

where the first inequality is due to triangle inequality, the second inequality is due to | ℓ ′ i | ≤ 1 , the third inequality is due to triangle inequality and the definition of neural networks, the forth inequality is due to parameter update rule (17) and Lemma B.1, and the fifth inequality is due to Condition 3.1.

Lemma C.15. Under the same conditions as Theorem 3.2, for all T 1 ≤ t ≤ T ∗ , we have max j,r,i | ρ ( t,b ) j,r,i | = max { O (√ log( mn/δ ) · σ 0 ( P -1) σ p √ d ) , O ( n √ log( n/δ ) log( T ∗ ) / √ d )} . Besides,

for all T 1 ≤ t ≤ T ∗ . Therefore, we can find an iterate with training loss smaller than 2 ϵ within T = T 1 + ⌊ ∥ W ( T 1 ) -W ∗ ∥ 2 F / ( ηϵ ) ⌋ = T 1 + ˜ O ( η -1 ϵ -1 mnd -1 σ -2 p ) iterations.

Proof of Lemma C.15. Note that max j,r,i | ρ ( t ) j,r,i | = max { O (√ log( mn/δ ) · σ 0 ( P -1) σ p √ d ) , O ( n √ log( n/δ ) log( T ∗ ) / √ d )} can be proved in the same way as Lemma C.9.

For any T 1 ≤ s ≤ T ∗ , by taking a summation of the inequality in Lemma C.13 and dividing ( s -T 1 ) H on both sides, we obtain that

According to the definition of T , we have

Then there exists an epoch T 1 ≤ t ≤ T ∗ such that

Thus, according to Lemma C.14, we have

Lemma C.16. Under the same conditions as Theorem 3.2, we have

for all j, j ′ ∈ {± 1 } , r, r ′ ∈ [ m ] and ( T 2 , 0) ≤ ( t, b ) ≤ ( T ∗ , 0) .

Proof. Now suppose that there exists (0 , 0) < ( ˜ T, 0) ≤ ( T ∗ , 0) such that ∑ n i =1 ρ ( t, 0) j,r,i /γ ( t,b ) j ′ ,r ′ = Θ(SNR -2 ) for all (0 , 0) < ( t, 0) < ( ˜ T, 0) . Then for ρ ( t,b ) j,r,i , according to Lemma C.1, we have

It follows that

Thus,

where the last equality is by the definition of S ( ˜ T -1) j,r as { i ∈ [ n ] : y i = j, ⟨ w ( ˜ T -1) j,r , ξ i ⟩ > 0 } ; the last inequality is by Lemma B.1 and the fifth statement of Lemma C.8. And

According to the third statement of Lemma C.8, we have max i ∈ S + | ℓ ′ ( ˜ T -1 ,b ( ˜ T -1) i ) i | ≤ C 2 min i ∈ S ( ˜ T -1 ,b ( ˜ T -1) i ) j,r | ℓ ′ ( ˜ T -1) i | . Then by combining (52) and (53), we have

On the other hand, we will now show ∑ n i =1 ρ ( t, 0) j,r,i γ ( t, 0) j ′ ,r ′ ≤ Θ(SNR -2 ) for t ≥ T 2 by induction. By Lemma B.1 and (52), we have

And, by (46), we know that at t = T 2 , we have

Suppose ∑ n i =1 ρ ( T, 0) j,r,i γ ( T, 0) j ′ ,r ′ ≤ Θ(SNR -2 ) . According to the decomposition, we have:

And we have that

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 ) j,r , ̂ y i µ ⟩ is persistent through out the epoch. Then, without loss of generality, we suppose ⟨ w ( T,b ) j,r , µ ⟩ > 0 . Thus, the update rule of γ is:

Therefore,

And, by (52), we have

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

where the last inequality is due to the induction hypothesis, the third statement of Lemma C.8, and Lemma B.5. Thus, by induction, we have for all T 1 ≤ t ≤ T ∗ that

And for ( T 1 , 0) ≤ ( t, b ) ≤ ( T ∗ , 0) , we can bound the ratio as follows:

where the first inequality is due to the update rule of ρ ( t,b ) j,r,i and ρ ( t,b ) j,r,i . Thus, we have completed the proof.

C.3 Test Error Analysis

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

C.3.1 Test Error Upper Bound

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

Theorem C.17 (Second part of Theorem 3.2) . Under the same conditions as Theorem 3.2, then there exists a large constant C 1 such that when n ∥ µ ∥ 2 2 ≥ C 1 ( P -1) 4 σ 4 p d , for time t defined in Lemma C.15, we have the test error

̸

where C 2 = O (1) .

Proof. The proof is similar to the proof of Theorem E.1 in Kou et al. (2023). The only difference is substituting ξ in their proof with ( P -1) ξ .

C.3.2 Test Error Lower Bound

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

Lemma C.18. For ( T 1 , 0) ≤ ( t, b ) < ( T ∗ , 0) , denote g ( ξ ) = ∑ j,r j ( P -1) σ ( ⟨ w ( t,b ) j,r , ξ ⟩ ) . There exists a fixed vector v with ∥ v ∥ 2 ≤ 0 . 06 σ p such that

for all ξ ∈ R d .

Proof of Lemma C.18. The proof is similar to the proof of Lemma 5.8 in Kou et al. (2023). The only difference is substituting ξ in their proof with ( P -1) ξ .

Lemma C.19 (Proposition 2.1 in Devroye et al. (2018)) . The TV distance between N (0 , σ 2 p I d ) and N ( v , σ 2 p I d ) is smaller than ∥ v ∥ 2 / 2 σ p .

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

Theorem C.20 (Third part of Theorem 3.2) . Suppose that n ∥ µ ∥ 4 2 ≤ C 3 d ( P -1) 4 σ 4 p , then we have that L 0 -1 D ( W ( t, 0) ) ≥ p +0 . 1 , where C 3 is an sufficiently large absolute constant.

Proof. The proof is similar to the proof of Theorem 4.3 in Kou et al. (2023). The only difference is substituting ξ in their proof with ( P -1) ξ .

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 ) i = ℓ ′ [ y i · f ( W ( t,b ) , x i )] , then the adversarial point of W ( t,b ) is W ( t,b ) + ̂ ϵ ( t,b ) , where

Then the training update rule of the parameter is

We will show that the noise term will be small if we train with the SAM algorithm. We consider the first stage where t ≤ T 1 where T 1 = mB/ (12 nη ∥ µ ∥ 2 2 ) . Then, the following property holds.

Proposition D.2. Under Assumption 3.1, for 0 ≤ t ≤ T 1 , we have that

Besides, γ ( T 1 , 0) j,r = Ω(1) .

Lemma D.3. Under Assumption 3.1, suppose (27) , (28) and (29) hold at iteration t . Then, for all r ∈ [ m ] , j ∈ {± 1 } and i ∈ [ n ] ,

̸

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

Lemma D.4. Under Assumption 3.1, suppose (61) , (62) , (63) hold at iteration t, b . Then, for all j ∈ {± 1 } and i ∈ [ n ] , F j ( W ( t,b ) j , x i ) ≤ 0 . 5 . Therefore -0 . 3 ≥ ℓ ′ i ≥ -0 . 7 .

̸

Proof. Notice that 1 / 12 < α , if the condition (61), (62), (63) holds, (27), (28) and (29) also holds. Therefore, by Lemma C.4, we know that for all j = y i and i ∈ [ n ] , F j ( W ( t,b ) j , x i ) ≤ 0 . 5 . Next we will show that for j = y i , F j ( W ( t,b ) j , x i ) ≤ 0 . 5 also holds.

According to Lemma D.3, we have

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

Since F j ( W ( t,b ) j , x i ) ∈ [0 , 0 . 5] we know that

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

Lemma D.5. Under Assumption 3.1, suppose (61) , (62) and (63) hold at iteration t, b . We have that if ⟨ w ( t,b ) j,r , ξ k ⟩ ≥ 0 , k ∈ I t,b and j = y k , then ⟨ w ( t,b ) j,r + ̂ ϵ ( t,b ) j,r , ξ k ⟩ < 0 .

Proof. We first prove that there for t ≤ T 1 , there exists a constant C 2 such that

Recall that we have

We have

and

leading to

From Lemma D.1, we have

̸

where we the last equality is by choosing τ = m √ B C 3 Pσ p √ d . Now we give an upper bound of ⟨ w ( t ) j,r , ξ k ⟩ , by (66) we have that

Combining (67) and (68) completes the proof.

Lemma D.6. Under Assumption 3.1, suppose (61) , (62) , (63) hold at iteration t, b . Then (62) also holds for t, b +1 .

Proof. Now consider the SAM algorithm. Recall that

Case 1: i / ∈ I t,b . In this case, clearly we have that ρ ( t,b +1) j,r,i = ρ ( t,b ) j,r,i ≤ 1 / 12 .

Case 2: i ∈ I t,b and ⟨ w ( t,b ) j,r , ξ i ⟩ ≥ 0 , then by Lemma D.5, we have that ⟨ w ( t ) j,r + ̂ ϵ ( t ) j,r , ξ k ⟩ < 0 , therefore we have that ρ ( t,b +1) j,r,i = ρ ( t,b ) j,r,i ≤ 1 / 12 .

Case 3: i ∈ I t,b and ⟨ w ( t,b ) j,r , ξ i ⟩ ≤ 0 then by (66) and triangle inequality, we can conclude that ρ ( t,b ) j,r,i can not reach a constant order,

Then we can give an upper bound for ρ ( t +1 ,b ) j,r,i since we only take one small step further,

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 1 such that the results in Proposition D.2 hold for all time (0 , 0) ≤ ( t, b ) ≤ ( ˜ T -1 , ˜ b -1) .We aim to prove that (61), (62), (63) also hold for iteration ( ˜ T -1 , ˜ b ) .

First, we prove that (61) holds for hold for iteration ( ˜ T -1 , ˜ b ) . Notice that

where the last inequality is by the fact that | ℓ ′ ( t,b +1) i | ≤ 1 and σ ′ ≤ 1 . Notice that ˜ T -1 ≤ T 1 , we can conclude that,

Second, by Lemma D.6, we know that (62) holds for ( ˜ T -1 , ˜ b ) .

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

and we have

and thus ⟨ w ( ˜ T -1 , ˜ b -1) j,r + ̂ ϵ ( ˜ T -1 , ˜ b -1) j,r , ξ i ⟩ < 0 which leads to

Therefore, we have

̸

Therefore, the induction is completed, and thus Proposition D.2 holds.

Next, we will prove that γ ( t ) j,r can achieve Ω(1) after T 1 = mB/ (12 nη ∥ µ ∥ 2 2 ) iterations. By Lemma B.6, we know that there exists c 3 · T 1 epochs such that at least c 4 · H batches in these epochs, satisfy

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

If ⟨ w ( t,b ) j,r + ̂ ϵ ( t,b ) j,r , µ ⟩ ≥ 0 , we have

If ⟨ w ( t,b ) j,r + ̂ ϵ ( t,b ) j,r , µ ⟩ < 0 , we have

Therefore, we have

Lemma D.7. Suppose Condition 3.1 holds. Then we have that ∥ ∥ w ( T 1 , 0) j,r ∥ ∥ 2 = Θ( σ 0 √ d ) and

Besides, for S ( t,b ) i and S ( t,b ) j,r defined in Lemma B.3 and B.4, we have that

D.2 Test Error Analysis

Proof of Theorem 4.1. Recall that

by triangle inequality we have

By the condition on σ 0 and Lemma B.2, we have

By taking the inner product with µ and ξ i , we can get

and

Then, we have

and

̸

where the last inequality is by the condition on σ 0 and γ ( T 1 , 0) j,r = Ω(1) . Thus, it follows that

By triangle inequality, we have

and

This leads to

̸

In addition, we also have for t ≤ T 1 and j = y i that

̸

Now let ¯ S (0 , 0) i denote { r : ⟨ w (0 , 0) y i ,r , ξ i ⟩ > σ 0 σ p √ d ⟩} and let ¯ S (0 , 0) j,r denote { i ∈ [ n ] : y i = j, ⟨ w ( t,b ) y i ,r , ξ i ⟩ > σ 0 σ p √ d } . By the condition on σ 0 , we have for t ≤ T 1 that

̸

for any r ∈ ¯ S (0 , 0) i or i ∈ ¯ S (0 , 0) j,r . Therefore, we have ¯ S (0 , 0) i ⊆ S ( T 1 , 0) i and ¯ S (0 , 0) j,r ⊆ S ( T 1 , 0) j,r and hence

where Φ( · ) is the CDF of the standard normal distribution.

Now we can give proof of Theorem 4.

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

where ˜ γ ( T 1 , 0) j,r = Θ(1) , ˜ ρ ( T 1 , 0) j,r,i ∈ [0 , 1 / 12] , ˜ ρ ( T 1 , 0) j,r,i ∈ [ -β -10 √ log(6 n 2 /δ ) /dn, 0] . Then the SGD start at W (0 , 0) := ˜ W ( T 1 , 0) . Notice that by Lemma D.7, we know that the initial weights of SGD (i.e., the end weight of SAM) W (0 , 0) still satisfies the conditions for Subsection C.1 and C.2. Therefore, following the same analysis in Subsection C.1 and C.2, we have that there exist t = ˜ O ( η -1 ϵ -1 mnd -1 P -2 σ -2 p ) such that L S ( W ( t, 0) ) ≤ ϵ . Besides,

for j ∈ [ ± 1] and r ∈ [ m ] where

Next, we will evaluate the test error for W ( t, 0) . Notice that we use ( t ) as the shorthand notation of ( t, 0) . For the sake of convenience, we use ( x , ̂ y, 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

̸

̸

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 ( x , ̂ y ) ∼D ( ̂ yf ( W ( t ) , x ) ≤ 0 ) . To achieve this, we write x = ( ̂ y µ , ξ ) , and get

The inner product with j = ̂ y can be bounded as

where the inequality is by Lemma B.1; the second equality is obtained by plugging in the coefficient orders we summarized at (70); the third equality is by the condition SNR = ∥ µ ∥ 2 /Pσ p √ d ; the fourth equality is due to ∥ µ ∥ 2 2 ≥ C · P 2 σ 2 p log( n/δ ) in Condition 3.1, so for sufficiently large constant C the equality holds; the last equality is by Lemma D.7. Moreover, we can deduce in a similar manner that

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

Denote g ( ξ ) as ∑ r σ ( ⟨ w ( t ) -̂ y,r , ξ ⟩ ) . According to Theorem 5.2.2 in Vershynin (2018), we know that for any x ≥ 0 it holds that

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

where the first inequality is by triangle inequality, the second inequality is by the property of ReLU; and the last inequality is by Cauchy-Schwartz inequality. Therefore, we have

and since ⟨ w ( t ) -̂ y,r , ξ ⟩ ∼ N ( 0 , ∥ w ( t ) -̂ y,r ∥ 2 2 σ 2 p ) , we can get

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

where for the first inequality, we used Lemma B.1; for the second inequality, we used the definition of ρ, ρ ; for the second to last equation, we plugged in coefficient orders. We can thus upper bound the 2 -norm of w ( t ) j,r as:

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

based on the coefficient order ∑ n i =1 ρ ( t ) j,r,i /γ ( t ) j,r = Θ(SNR -2 ) , the definition SNR = ∥ µ ∥ 2 / ( σ p √ d ) , and the condition for d in Condition 3.1; and also ∥ w (0) j,r ∥ 2 = Θ( σ 0 √ d ) based on Lemma D.7. With this and (73), we analyze the key component in (80):

It directly follows that

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

where the second inequality is by (79) and plugging (76) into (75), the third inequality is due to the fact that ( s -t ) 2 ≥ s 2 / 2 -t 2 , ∀ s, t ≥ 0 .

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

where C = O (1) , the last inequality holds since σ 2 0 ≤ 0 . 5 Cd -1 log(1 /ϵ ) and d ≥ 2 C -1 P -1 σ -2 p nα 2 log(1 /ϵ ) .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- **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\nWEAKNESSES:\n- **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.\n\nQUESTIONS:\nCould you please share your thoughts on the potential behavior of SAM within the gray region depicted in Figure 1?\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- Discussion of the results seems not quite comprehensive; see the comments below.\n\nQUESTIONS:\n- 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)\n\nLIMITATIONS:\nPlease see the comments in the Question section.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nSee weakness.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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”.\n\nQUESTIONS:\n1) 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?\n\nLIMITATIONS:\nSee above regarding mathematical assumptions. Societal impact is not addressed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}