Benign Overfitting and Grokking in ReLU Networks for XOR Cluster Data
Authors: Zhiwei Xu
Yutong Wang
Spencer Frei
Gal Vardi
Wei Hu
Venue: ICLR 2024 poster (2024)
Keywords: benign overfitting
grokking
neural networks
feature learning
interpolation
theory
Primary area: learning theory

Abstract: Neural networks trained by gradient descent (GD) have exhibited a number of surprising generalization behaviors. First, they can achieve a perfect fit to noisy training data and still generalize near-optimally, showing that overfitting can sometimes be benign. Second, they can undergo a period of classical, harmful overfitting---achieving a perfect fit to training data with near-random performance on test data---before transitioning (''grokking'') to near-optimal generalization later in training. In this work, we show that both of these phenomena provably occur in two-layer ReLU networks trained by GD on XOR cluster data where a constant fraction of the training labels are flipped. In this setting, we show that after the first step of GD, the network achieves 100\% training accuracy, perfectly fitting the noisy labels in the training data, but achieves near-random test accuracy. At a later training step, the network achieves near-optimal test accuracy while still fitting the random labels in the training data, exhibiting a ''grokking'' phenomenon. This provides the first theoretical result of benign overfitting in neural network classification when the data distribution is not linearly separable. Our proofs rely on analyzing the feature learning process under GD, which reveals that the network implements a non-generalizable linear classifier after one step and gradually learns generalizable features in later steps.

Published as a conference paper at ICLR 2024
BENIGN OVERFITTING AND GROKKING IN RELU NETWORKS FOR XOR CLUSTER DATA
Zhiwei Xul, Yutong Wang Spencer Frei', Gal Vardi Wei Hu
University of Michigan, UUiiversity of California, Davis, TTT--Cicago and Hebrew University {zhiweixu, yutengw, vvh whh]eumich ch. edu, sfrei Bucdavis s edu, salvardi@ttic.... edu
ABSTRACT
Neural networks trained by gradient descent (GD) have exhibited a number of sur- prising generalization behaviors. First, they can achieve a perfect fit noisy train- ing data and still generalize near-optimally, showing that overfitting can sometimes be benign. Second, they can undergo a period of classical, harmful overfitting achieving a perfect fit to training data with near-random performance on test data-before transitioning ("grokking") to near-optimal generalization later training. In this work, we show that both of these phenomena provably occur in two-layer ReLU networks trained by GD on XOR cluster data where a constant fraction of the training labels are flipped. In this setting, show that after the first step of GD, the network achieves 100% training accuracy, perfectly fitting the noisy labels in the training data, but achieves near-random test accuracy. At a later training step, the network achieves near-optimal test accuracy while still fitting the random labels in the training data, exhibiting "grokking" phenomenon. This provides the first theoretical result of benign overfitting in neural network classification when the data distribution is not linearly separable. Our proofs rely analyzing the feature learning process under GD, which reveals that the network implements a non-genere non-generalizable linear classifier after one step and gradually learns generalizable features in later steps.
INTRODUCTION
Classical wisdom in machine learning regards overfitting to noisy training data harmful for generalization, and regularization techniques such as early stopping have been developed to prevent overfitting. However, modern neural networks can exhibit number of counterintuitive phenomena that contravene this classical wisdom. Two intriguing phenomena that have attracted significant attention in recent years are benign overfitting (Bartlett al., 2020) and grokking (Power et al., 2022):
Benign overfitting: model perfectly fits noisily labeled training data, but still achieves near-optimal test error.
Grokking: A model initially achieves perfect training accuracy but no generalization (i.e. no better than random predictor), and upon further training, transitions to almost perfect generalization.
Recent theoretical work has established benign overfitting in a variety of settings, including linear regression (Hastie et al., 2019; Bartlett et al., 2020), linear classification (Chatterji & Long, 2021 la;
Wang & Thrampoulidis, 2021), kernel methods (Belkin et al., 2019; Liang & Rakhlin, 2020), and neural network classification (Frei et al., 2022b; Kou al., 2023). However, existing results of benign overfitting in neural network classification settings distributions, leaving open the question of how benign overfitting can occur in fully non-linear settings. For grokking, several recent papers (Nanda et al., 2023; Gromov, 2023; Varma al., 2023) have proposed explanations, but to the best of our knowledge, no prior work has established rigorous proof of grokking in neural network setting.
restricted to linearly separable data
In this work, we characterize a setting in which both benign overfitting and grokking provably occur. We consider two-layer ReL.U network trained by gradient descent binary classification task

---

Published as a conference paper at ICLR 2024
1.0
1.0
accuracy
train (med)
0.8
(med)
0.8
20-80
DCC aCC-0.3
10
10a
10:
10a
10a
steps
steps
Figure I: Comparing train and test accuracies of two-layer neural network (2.1) trained noisily labeled XOR data over 100 independent runs. Left/right panel shows benign overfitting and grokking when the step size is larger/smaller compared the weight initialization scale. For plotting the x-axis, we add to time so that the initialization & = 0 can be shown in log scale. See Appendix A.7 for details of the experimental setup.
Cluster detail
Step
Step
Step
Neurons
Neurons @ Step
Figure 2: Left four panels: 2-dimensional projection the noisily labeled XOR cluster data (Defini- tion 2.1) and the decision boundary of the neural network (2.1) classifier restricted to the subspace spanned by the cluster means times the neuron weights plotted at times / = and 15.
and 15. Right two panels: 2-dimensional projection of
defined by an XOR cluster data distribution (Figure 2). Specifically, datapoints from the positive class drawn from mixture of two high-dimensional Gaussian distributions J.N () f (- I), and datapoints from the negative class are drawn from 1 N r) + (-- 7), where (4 and μl2 orthogonal vectors. We then allow constant fraction of the labels be flipped. In this setting, we rigorously prove the following results: (i) One-step catastrophic overfitting: After one gradient descent step, the network perfectly fits every single training datapoint (no matter if it has clean flipped label), but has test accuracy close to 50%, performing no better than random guessing. (ii) Grokking and benign overfitting: After training for steps, the network undergoes "grokking" period from catastrophic to benign overfitting- eventually reaches near 100% test accuracy, while maintaining 100% training accuracy the whole time. This behavior can be seen in Figure where we also see that with a smaller step size the same grokking phenomenon occurs but with delayed time for both overfitting and generalization.
Our results provide the first theoretical characterization of benign overfitting in a truly non- non-linear setting involving training neural network on non-linearly separable distribution. Interestingly, prior work on benign overfitting in neural networks for linearly separable distributions (Frei et al., 2022b; Cao et al., 2022: Xu Gu, 2023; Kou al., 2023) have not shown time separation between catastrophic overfitting and generalization, which suggests that the XOR cluster data setting is fundamentally different.
Our proofs rely on analyzing the feature learning behavior of individual neurons over the gradient descent trajectory. After one training step, we prove that the network approximately implements a linear classifier over the underlying data distribution, which is able to overfit all the training datapoints but unable to generalize. Upon further training, the neurons gradually align with the core features =t1 and =22, which is sufficient for generalization. See Figure for visualizations of the network's decision boundary and neuron weights at different time steps, which confirm our theory.
1.1 ADDITONAL RELATED WORK
Benign overfitting. The literature on benign overfitting (also known as harmless interpolation) now immense; for general overview, refer the readers to the surveys Bartlett et al. (2021); Belkin (2021); Dar et al. (2021). We focus here on those works on benign overfitting in neural networks. Frei et al. (2022b) showed that two-layer networks with smooth leaky ReLU activations
2

---

Published as a conference paper at ICLR 2024
trained by gradient descent (GD) exhibit benign overfitting when trained high-dimensional binary cluster distribution. Xu & Gu (2023) extended their results to more general activations like ReLU. Cao et al. (2022) showed that two-layer convolutional networks with polynomial-ReLJ activations trained by GD exhibit benign overfitting for image-patch data; Kou et al. (2023) extended their results to allow for lahel-flipiing noise and standard ReI.U activations. Each of these works used a trajectory-basee analysis and none of them identified a grokking phenomenon. Frei et al. (2023a); Kornowski al. (2023) showed how stationary points of aangin-maximization problems associated with homogeneous neural network training problems can exhibit benign overfitting. Finally, Mallinar et al. (2022) proposed a taxonomy of overfitting behaviors neural networks. whereby overfitting is "catastrophic" if test-time performance is comparable to a random guess, "benign" if it is near-optimal, and "tempered" it lies between catastrophic and benign.
Grokking. The phenomenon of grokking was first identified by Power et al. (2022) in decoder- only transformers trained on algorithmic datasets. Liu et al. (2022) provided an effective theory of representation learning to understand grokking. Thilak et al. (2022) attributed grokking to the slingshot mechanism, which can be measured by the cyclic phase transitions between stable and unstable training regimes. Zunkovič & Ilievski (2022) showed time separation between achieving zero training error and zero test error in a binary classification task on a linearly separable distribution. Liu et al. (2023) identified large initialization scale together with weight decay as mechanism for grokking. Barak et al. (2022); Nanda et al. (2023) proposed progress metrics to measure the progress towards generalization during training. Davies et al. (2023) hypothesized a pattern-learning model for grokking and first reported model-wise grokking phenomenon. Merrill et al. (2023) studied the learning dynamics in a two-layer neural network on a sparse parity task, attributing grokking the competition between dense and sparse subnetworks. Varma et al. (2023) utilized circuit efficiency to interpret grokking and discovered two novel phenomena called ungrokking and -rookinn..
Feature learning for XOR distributions. The behavior of neural networks trained on the XOR cluster distribution we consider here, its variants like the sparse parity problem, have been extensively studied in recent years. Wei et al. (2019) showed that neural networks in the mean-field regime, where neural networks can learn features, have better sample complexity guarantees than neural networks in the neural tangent kernel (NTK) regime in this setting. Barak et al. (2022); Telgarsky (2023) examined the sample complexity of learning sparse parities on the hypercube for neural networks trained by SGD. Most related to this work, Frei et al. (2022a) characterized the dynamics of GD in ReLU networks in the same distributional setting we consider here, namely the XOR cluster with label-flippigg noise. They showed that by early-stopping, the neural network achieves perfect (clean) test accuracy although the training error close to the label noise rate; in particular, their network achieved optimal generalization without overfitting, which is fundamentally different from our result. By contrast, we show that the network first exhibits catastrophic overfitting before transitioning to benign overfitting later in training."
2
PRELIMINARIES
2.1 NOTATION
For a vector x, denote its Euclidean norm by ||-]| For a matrix X, denote its Frobenius norm by || and its spectral norm by || Denote the indicator function by Il(-). Denote the sign of scalar x by sgn(x). Denote the cosine similarity of two vectors u, e by cossim(u, v) == [[]|||]] Denote a multivariate Gaussian distribution with mean vector and covariance matrix Σ by N(u, Σ). Denote by Σ, N, Σ,) a mixture of Gaussian distributions, namely, with probability qj- the sample is generated from N(,, Σz). Let Ip be the x identity matrix. For a finite set A = {a denote the uniform distribution on A by UnifA. For a random variable X, denote its expectation by ] For an integer d 2 1, denote the set {1, ... d} by [d]. For finite set A, let [A] be its cardinality. We use {+p represent the {+14, - w.) For two positive sequences ( ( (vr.). we say In = Og.)) (respectively I, = )), if there exists a universal constant C > 0 such
fu
'The reason for the different behaviors hetween our work and Frei et al. (2022a) setting with larger signal-to-noise ratio (i.c., the norm of the cluster means is larger than the one we consider).
is because they work in
3

---

Published as a conference paper at ICLR 2024
that 2n I Cyn (respectively Σn 2 Cya) for all n, and say 2n o(vr) if nim Bn. In (yn,) I, = (yn) and yn (zx)).
0. We say
2.2 DATA GENERATION SETTING
Let μr, μ22 E R be two orthogonal vectors, i.e. vi P2 0.2 Let E [0, 1/2) be the label flipping probability.
Definition 2.1 (XOR cluster data). Define Pelemn the distribution over the space RR x {+1} of labelled data such that datapoint (z, D) ~ Palean is generated according to the following procedure:
First, sample the label Uni Second, generate follows:
(1) If = 1, then [+ + 2 11. I1)1p)) Ip);
(2) y = 1, then x M+ 1p) + 2 N 12, Ip).
Define to be the distribution over Rº x [=1} which is the y-noise-coruuted version of Polcan- namely: generate sample (xr, y) P. first generate (x. V) Pelemn and then let y ÿ with
probability 9), and = ý with probability y/.
We consider n training datapoints ((r,,)))) y.))== the sample size to be sufficiently large (i.e., larger than any universal constant appearing in this paper). Note the x; 's are from a mixture of four Gaussians centered at =11 and + We denote centers: = { [ for convenience. For simplicity, assume || || omit the subscripts and denote them by ||/]]-
generated i.i.d from the distribution P. We assume
2.3 NEURAL NETWORK, Loss FUNCTION, AND TRAINING PROCEDURE
We consider two-layer neural network of width m of the form
m
f(x; W)
Σ a (w),,
(2.1)
where Wi, ....... Wm E R2 are the first-layer weights, aI..... am E IR are the second-layer weights, and the activation 6(2) == max =} is the ReLU function. We denote W [wi,..., wmJ IR and a = [a....., am] E I=.. We assume the accond-laye weights are sampled according to aj Lid Uniff } and are fixed during the training process.
We define the empirical risk using the logistic loss function ((2) = log(1 + exp(-2)) Ż(W)= =
E (y (x,; W)). We use gradient descent (GD) ((+1) W(() V (W(t) to update the first-layer weight matrix W, where a is the step size. Specifically, at time we randomly initialize the weights by
(0) ltd. (0, Ip). j j [m].
where wit is the initialization variance; at each time step t 0,1,2..... the GD update can be calculated as
o((w()) aa,
w (++)) w/) aa
9 o' d'((w)))
,xi)), E [m]. (2.2)
Ow, mml
where 1:2h ;= -'(ysl ((i;)()))).
MAIN RESULTS
Given large enough universal constant C, we make the following assumptions:
Our results hold when /41 and #2 near-orthogonal. We assume exact orthogonality for ease of presenta-
tion.
4

---

Published as a conference paper at ICLR 2024
(A1) The norm of the mean satisfies ||a|| 2 Cn vp.
(A2) The dimension of the feature space satisfies p 2 Cn ||ll?.
(A3) The noise rate satisfies 1) 1 1/C.
(A4) The step size satisfies S 1/(Cnp).
(A5) The initialization variance satisfies Wnit nm3 S a ar||a|1
(A6) The number of neurons satisfies m 2 0.02
Assumption (A1) concerns the signal-to-noise ratio (SNR) in the distribution, where the order 0.51 can be extended to any constant strictly larger than 2. The assumption of high-dimensionality (A2) is important for enabling benign overfitting, and implies that the training datapoints are near-orthogonal. For given n, these two assumptions are simultancously satisfied if |||||| ((p²) where / E (4, 2) and is a sufficiently large polynomial in n. Assumption (A3) ensures that the label noise rate is at most a constant. While Assumption (A4) ensures the step size is small enough to allow for a variant of smoothness between different steps, Assumption (A5) ensures that the step size is large relative to the initialization scale so that the behavior of the network after a single step of GD significantly different from that at random initialization. Assumption (A6) ensures the number of neurons is large enough to allow for concentration arguments at random initialization.
With these assumptions in place, can state our main theorem which characterizes the training error and test error of the neural network at different times during the training trajectory.
Theorem 3.1. Suppose that Assumptions (A1)-(A6) hold. With probability at least 1 R(1) O(1/ mm) over the random data generation and initialization of the weights. we have:
The classifier of ((::)( W( can correctly classify all training datapoints for sts Vn:
y = sgu(f(x; W())), Vi [n].
The classifier sgn(s(1; W()) has near-random test error at f. 1:
/ n - < ((z,)))PPn (v & sgu(f( W()))) $ 21
TThe classifier sgn (f W(1) generalizes when 0.01 St Vn:
PP ((..))) (y + (g W(+)) S exp( 0.99 /p)) exp( exp(-S((n201)
Theorem 3.1 shows that at time f = 1, the network achieves 100% training accuracy despite the constant fraction of flipped labels in the training data. The second part of the theorem shows that this overfitting is catastrophic as the test eтror is close to that of a random guess, On the other hand, by the first and third parts of the theorem, long as the time step t satisfies 0.000 St yn, the network continues to overfit to the training data while simultaneously achieving test error exp( (nn22011)) )), which guarantees near-zero test error for large n. particular, the network exhibits benign overfitting, and it achieves this by grokking. Notably, Theorem 3.1 is the first guarantee for benign overfitting in neural network classification for a nonlinear data distribution, in contrast to prior works which required linearly separable distributions (Frei al., 2022b; 2023a; Cao et al., 2022; Xu & Gu, 2023; Kou et al., 2023; Kornowski et al., 2023).
We note that Theorem 3.1 requires upper bound on the number of iterations of gradient descent, Le. it does not provide a guarantee as t > oc. At a technical level, this is needed so that we can
guarantee that the ratio of the sigmoid losses between all samples r(t) := je( Si is close to
g)
1, and we show that this holds if < vn. This property prevents the training data with flipped labels from having an out-sized influence on the feature learning dynamics. Prior works in other settings have shown that rtt) is at most a large constant for any step t for a similar purpose (Frei et al., 2022b; Xu & Gu, 2023), however the dynamics of learning in the XOR setting are more intricate and require a tighter bound on r rt). We leave the question of generalizing our results to longer training times for future work. In Section 4, we provide overview of the key ingredients to the proof of Theorem 3.1.
5

---

Published as a conference paper at ICLR 2024
4 PROOF SKETCH
We first introduce some additional notation. For i E [n]] let Σ; E centers = { ) be the mean of the Gaussian from which the sample (r,, Mr) is drawn. For each v E centers, define In {i E [n] I; v), i.e., the set of indices such that æ; belongs to the cluster centered at v. Thus, {Iv recenters is a partition of [n]. Moreover, define = {i E [r] y = Yi } and [n] y o V to be the set of clean and noisy samples, respectively, Further define for each y E centers the following sets:
= {i =
C, CnI, and N=
Let Cv [Cul and ny Wel. Define the training input data matrix X [2] , F.]] 8 E (0, (0,10-8/4) /4) be a universal constant.
Let
In Section 4.1, we present several properties satisfied with high probability by the training data and random initialization, which are crucial in our proof. In Section 4.2, we outline the major steps in the proof of Theorem 3.1.
4.1 PROPERTIES OF THE TRAINING DATA AND RANDOM INITIALIZATION
Lemma 4.1 (Properties of training data). Suppose Assumptions (AI) and (A2) hold. Let the training data {<),,:)) be sampled i.i.d from P as in Desinition 2.1. With probability at least 1 O(n the training data satisfy properties (BI)-(B4) defined below.
(B1) For all k E [n]].
max (It -ū,, v) < 10v/oggnlle and ||| 1/ppl/=p-11llppp S 10v plogooonn n. VEcenters
(B2) For each i, [n] such that i + k we have (1i, (T,, Fx>i S 1og.log..
(B3) For E centers, have lov +n n/4/ S Ven Venlogn and (no- (ec +nu) S Venn log
(B4) For v E centers, we have lcv + fyy C-v n-vl 2 n 1/2-- and mn n- 2 1p7 12--.
Denote by Gdata the of training data satisfying conditions (B1)-(B4). Thus, the result can be stated succinctly as TPXX E Gdiaa) 2 Onn-))
The proof of Lemma 4.1 can be found in Appendix A.2.1. Conditions (B1) and (B2) are essentially the same Frei et al. (2022b, Lemma 4.3) Chatterji & Long (2021b, Lemma 10). Conditions (B3) and (B4) concern the number of clean and noisy examples in each cluster, and can proved by concentration and anti-concentration arguments, respectively.
Lemma 4.1 has an important corollary.
Corollary 4,2 (Near-orthogonality of training data). Suppose Assumptions (AI), (A2), and Conditions (BI), (B2) from Lemma 4. I all hold. Then for all i f k E n,
|cossim(a;; xxmI <
C'2²
This craoooothogonallty comes from the high dimensionality of the feature space (i.c., Assump- tion (A2)) and will be crucially used throughout the proofs optimization and generalization of the network. The proof of Corollary 4.2 can be found Appendix 1.2.1.
Next, divide the neuron indices into sets according to the sign of the corresponding second-
layer weight:
Jrea Ij E [m] : ag 0); Jnag = (i [m] aj <0.
We will conveniently call them positive and negative neurons. Our next lemma shows that some properties of the random initialization hold with a large probability. The proof details can be found in Appendix A.3.1.
Lemma 4.3 (Properties of the random weight initialization). Suppose Assumptions (A1), (A2) and (A6) hold. All conditions below simultaneously hold with probability at least 1 O(n -)) over the random initialization:

---

Published as a conference paper at ICLR 2024
(CI) |W(00]] S 3 mp.
(C2) [Jpos| 2 m/3 and [Treg] 2 m/3.
The result can be stated equivalently as follows: Denote the set of W satisfying condition (CI) by Gw. Denote the set of a = (az)ā=1 satisfying condition (C2) by Ga. Then IP(a E Ga, W(
Gw) 2 1 - ((n
We say that the sample activates neuron j at time if (wy .4i) > O. Now, for each neuron j [m].
time t 2 0 and E centers, define the set of indices of samples 2; with clean (resp. noisy) labels from the cluster centered that activates neuron j at time
(1) vi) > 0} (resp. N == N (wf aii 0}). C ii Cv (w)
(4.1)
Moreover, we define
49 = 1c -1 and [0 d d
For R [0, 1/2) and v centers, a neuron j is said to be (v, -allgnndd if
Do 1/-- and max(d_ d min( c-u) 2(n + n- Vn
(4.2)
The first condition ensures that at initialization, there are at least n2/22xx many more samples from cluster u activating the j-th neuron than from cluster - after accounting for cancellations from the noisy labels. The second is technical condition necessary for trajectory analysis. A neuron j is said to be (v, x)aalineed if it either (v, -llinned or (-v, al)-aieeedd
Lemma 4.4 (Properties of the interaction between training data and initial weights), Suppose Assumptions (AI)-(A3) and (A6) hold. Given a E GA (defined in Lemma 4.3) and X E Idau (defined in Lemma 4.1), the followings hold with probability at least 1 O(" -5) over the random initialization (
(DI) For all i [n], the sample x; activates a large proportion of positive and negative neurons, i.e., Kij E Toos : (w]) (0) xi) > OH 2 m/7 7 and K E Ixog (w (0) hold.
xz) > 0H 2 m/7 both
(D2) For all centers and K. E [0, }), both Hjj Jrea : is (u, -aalign))) 2 mn 10e 10s and Hj E Jnag ; is (v, k))aligned]] 2 mn 10e,
(D3) For all E centers, we have Hj E Jrea ;j is (#v, ) 2 (1 10n 1000000)) Moreover, the same statement holds if "Jros" is replaced with "Jneg" everywhere.
20c [Jrool
(D4) For all v E centers and =[0,2) let Júpos == (i E Joes : is (v, allignd),.. Then EjeJo (c. -c, - d ) to I Moreover, the same statement holds if " JTos" is replaced with "Jxeg" " everywhere.
Condition (D1) makes sure that the neurons spread uniformly initialization so that each datapoint activates at least constant fraction of positive and negative neurons. Condition (D2) guarantees that for each E centers, there are a fraction of neurons aligning with v more than -v. Condition (D3) shows that most neurons will somewhat align with either concentration result. For proof details, see Appendix A.3.2.
Condition (D4) is technical
Define the set Ggood
Ggoond == {aa (()) X): a a GA, X E G ((0) E Gw and conditions (DI)-(D4) hold),
whose probability is lower bounded by P((a, W()) [) E Ggood) (n--) This is conse-
quence of Lemmas 4.1, 4.3 and 4.4 (see Appendix A.3.3).
Definition 4.5. If the training data X and the initialization a, (0))) belong to Ggood we define this circumstance "good run."
7

---

Published as a conference paper at ICLR 2024
t
Number of neurons (Frequency)
0.02
0.02
0.02
0.02
0.03
0
0.05
00.10
0.20
0.02
0.02
(wf2) (4)
-1
-5
0
5
-2
2
2.5
0.0
2.5
-1
1e-13
1e-12
1e-11
16-11
1e-10
(wf?) (b)
Figure Histograms of inner products between positive neurons and μl or μ2 pooled over 100 independent runs under the same setting as in Figure 1. Top (resp. bottom) row: Inner products between positive neurons and ls (resp. (12). While the distributions of the projections of positive neurons w/ onto the μ1 and directions are nearly the same at times t = 0. 1, they become significantly more aligned with IP1 over time. See Appendix A.7 for details of the experimental setup.
4.2 PROOF SKETCH FOR THEBREM 3.1
In order for the network to learn a generalizable solution for the XOR cluster distribution, we would like positive neurons' (i.e., those with 0j > 0) weights Wy to align with =Ø1, and negative neurons' weights align with +22; we prove that this is satisfied for / [.o0001 vn). However, for we show that the network only approximates a linear classifier, which can fit the training data in high dimension but has trivial test error. Figure 3 plots the evolution of the distribution of positive neurons' projections onto both and μ2. confirming that these neurons are much more aligned with =un at later training time, while they cannot distinguish =U1 and +P2 at
Below we give a sketch of the proofs, and details are in Appendix A.5.
4.2.1 ONE-STEP CATASTROPHIC OVERFITTING
Under a good run, we have the following approximation for each neuron after the first iteration:
I(( (0) .r)) O)ger,, j e [m].
For details of this approximation, Appendix A.4.
Let 8tj = T(((w))) x.) 0). Then, for sufficiently large m, we can approximate the neural network
output at / = 1 as
m
a (w))))
xx))
~
a afola
8178111, x))
2n
1=1
j=1
(4.3)
a 8n
Missllias x)
yet x).
4n
This convergence is ensured by the strong law of large numbers, given the independence of the first-layer and second-layer weights at initialization. This implies that the neural network classifier (=;( (())) behaves similarly to the linear classifier sgn((>) Mer,, +)). It can be shown that this linear classifier achieves 100% training accuracy whenever the training data are near orthogonal (Frei al.. 2023b, Appendix D), but because each class has two clusters with opposing means, linear classifiers only achieve 50% test error for the XOR cluster distribution. Thus at time
I 1, the network is able fit the training data but is not capable of generalizing.
8

---

Published as a conference paper at ICLR 2024
4.2.2 MULTI-STEP GENERALIZATION
Next, we show that positive (resp. negative) neurons gradually align with one of =uj (resp. =22), and forget both of 4u2 (resp. =11), making the network generalizable. Taking the direction + as an example, we define sets of neurons
Ji ij E Joa j ; J2 {j Jueg ;j is (0e)
We have by conditions (D2)-(D3) of Lemma 4.4 that under good run,
F 2 mm 10c
1 2 (1 - 10n 20c) Tnogh
which implies that Ji contains certain proportion of Jros and J2 covers most of Jnag- The next lemma shows that neurons in J1 will keep aligning with +/11, but neurons in T2 will gradually forget +P1.
Lemma 4.6. Suppose that Assumptions (A1)-(A6) hold. Under a good run, we have that for 1 < t S Vn.
(allee1
+11) = 52
vm
JEJ
a||uu|?
a
al|p]l2 V log(n), Vmn
pl)1 =
J22
mm
JEJ
We can see that when t is large. Σ JeJs Kw ,1))1/111111 = o(jje,, (w))) +ba)/(FIl), thus for I ~ (++un, Ip), neurons with j E TI will dominate the output of f(z; W()). For the other three clusters centered -11 +/22, -02 generalization. Formally, we have the following theorem on generalization.
have similar results, which then lead the model to
Theorem 4.7. Suppose that Assumptions (A ((A6) hold. Under a good run, for On SLS Vn. the generalization error of classifier sgn( (f(r. W())) has an upper bound
P(x.))) (y # sgn(f(x; ((()) S exp (-2 (n' 1-200 ||0||
5 DISCUSSION
We have shown that two-layer neural networks trained on XOR cluster data with random label noise by GD reveal a number of interesting phenomena. First, early in training, the network interpolates all of the training data but fails to generalize to test data better than random chance, displaying a familiar form of (catastrophic) overfitting. Later in training, the network continues to achieve a perfect fit to the noisy training data but groks useful features that it achieve near-zero error on test data, thus exhibiting both grokking and benign overfitting simultaneously. Notably, this provides an example of benign overfitting in neural network classification for distribution which is not linearly separable.
In contrast to prior works on grokking which found the usage of weight decay to be crucial for grokking (Liu al., 2022: 2023), we observe grokking without any explicit forms of regularization, revealing the significance of the implicit regularization of GD. In our setting, the catastrophic overfitting stage of grokking occurs because early training. the network behaves similarly linear classifier. This linear classifier is capable of fitting the training data due to the high- dige-iimensionality of the feature space but fails to generalize as linear classifiers are not complex enough to achieve test performance above random chance for the XOR cluster. Later in training, the network groks useful features, corresponding to the cluster means, which allow for good generalization.
There are a few natural questions for future research. First, our analysis requires an upper bound on the number of training steps due to technical reasons; it is intriguing to understand the generalization behavior as time grows to infinity. Second, training data are orthoonl which requires that the ambient dimension is large relative to the number of samples. Prior work has shown with experiments that overfitting is less benign in this setting when the dimension is small relative to the number of samples (Frei et al., 2022a, Fig. 2); a precise characterization of the effect of high-dime data generalization remains open.
proof crucially relies upon the assumption that the

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
YW acknowledges support from the Eric and Wendy Schmidt AI in Science Postdoctoral Fellowship. a Schmidt Futures program. GV acknowledges support from the NSF and the Simons Foundation via the Collaboration on the Theoretical Foundations of Deep Learning. WH acknowledges support from the Google Research Scholar program.
REFERENCES
Boaz Barak, Benjamin L. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress deep learning: Sgd learns parities near the computational limit. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 117(49):30033-3007) 20200. 2020.
Peter L. Bartlett, Andrea Montanari, and Alexander Rakhlin. Deep learning: statistical viewpoint.
Acta Numerica, 30:87-201, 2021.
Mikhail Belkin. Fit without fear: remarkable mathematical phenomena of deep learning through the prism of interpolation. Acta Numerica, 30, 2021.
Mikhail Belkin, Alexander Rakhlin, and Alexandre B. Tsybakov. Does data interpolation contra- dict statistical optimality? In International Conference on Artificial Intelligence and Statistics (AISTATS), 2019.
Yuan Cao, Zixiang Chen, Mikhail Belkin, and Quanquan Gu. Benign overfitting in two-layer convolutional neural networks. arXiv preprint arXiv:2202.06526, 2022.
Niladri S. Chatterji and Philip M. Long. Finite-sample analysis of interpolating linear classifiers in the overparameterized regime. Journal of Machine Learning Research, 22(129):1-30, 2021a.
Niladri S. Chatterji and Philip M. Long. Finite-sample analysis of interpolating linear classifiers in the overparameterized regime. Journal of Machine Learning Research, 22 12312))1-30, 2021b. URL http: jmlr. org/
otg/papers papers v2222002222222 974 474. hhtll.
Yehuda Dar, Vidya Muthukumar, and Richard G. Baraniuk. A farewell to the bias-variance tradeoff? an overview of the theory overparameterized machine learning. Preprint, arXiv:2109.02::
2021.
Xander Davies, Lauro Langosco, and David Krueger. Unifying grokking and double descent. 2023.
Rick Durrett. Probability: theory and examples, volume 49. Cambridge university press, 2019.
Spencer Frei, Niladri S Chatterji, and Peter L Bartlett. Random feature amplification: Feature learning and generalization in neural networks. Preprint, arXiv:2202.07626 2022a.
Spencer Frei, Niladri Chatterji, and Peter Bartlett. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In Conference on Learning Theory (COLT), 2022b.
Spencer Frei, Gal Vardi, Peter L. Bartlett, and Nathan Srebro. Benign overfitting in linear classifiers and leaky relu networks from kkt conditions for margin maximization. In Conference on Learning Theory (COLT), 2023a.
Spencer Frei, Gal Vardi, Peter L. Bartlett, Nathan Srebro, and Wei Hu. Implicit bias in leaky relu net- works trained on high-dimensional data. In International Conference on Learning Representations, 2023b.
Andrey Gromov. Grokking modular arithmetic. Preprint, arXiv:2301.02677 2023.
Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan Tibshirani. Surprises in high- dimensional ridgeless least squares interpolation. arXiv preprint arXiv: iate:0903000000 2019.
10

---

Published as a conference paper at ICLR 2024
Guy Komnwski, Gilad Yehudai, and Ohad Shamir. From tempered benign overfitting in relu neural networks. Preprint, arXiv:2305.15144, 2023.
Yiwen Kou, Zixiang Chen, Yuanzhou Chen, and Quanquan Gu. Benign overfitting for two-layer relu convolutional networks. In International Conference on Machine Learning (ICML), 2023.
Tengyuan Liang and Alexander Rakhlin. Just interpolate: Kernel "ridgeless" regression can generalize. Annals of Statistics, 2020.
Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud. Max Tegmark, and Mike Williams. To- wards understanding grokking: An effective theory of representation learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neu- ral Information Processing Systems, volume 35, pp. 34651-34663. Curran Associates, Inc., 2022. URL file #tle//ttcsceess192ceeeeeeddcccee
neurips pacppaper_fiiie www.ffellsppppee/....
Paper Conference pdf.
Ziming Liu, Eric J. Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. In International Conference on Learning Representations (ICLR), 2023.
Neil Mallinar, James B Simon, Amirhesam Abedsoltan, Parthe Pandit, Mikhail Belkin, and Preetum Nakkiran. Benign, tempered, or catastrophic: A taxonomy of overfitting. In Advances in Neural Information Procesisng Systems (NeurIPS), 2022.
William Merrill, Nikolaos Tsilivis, and Aman Shukla. A tale of two circuits: Grokking as competition of sparse and dense subnetworks. 2023.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretabilityy Preprint, arXiv:2301.05217, 2023.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: General- ization beyond overfitting on small algorithmic datasets. Preprint, arXiv:2201.02177, 2022.
Matus Telgarsky. Feature selection and low test error in shallow low-rotation reLU networks. In International Conference Learning Representations (ICLR), 2023.
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. 2022.
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. Preprint, arXiv:2309.023300 2023.
Martin J. Wainwright. Digeesimnssi Statistics: A Non-Asymptotic Viewpoint. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2019. doi:
10.1017/9781 108627771.
Ke Wang and Christos Thrampoulidis. Binary classification of gaussian mixtures: Abundance of support vectors, benign overfitting and regularization. Preprint, arXiv:201 1.00144,, 2021.
Colin Wei, Jason D. Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets v.s. their induced kernel. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Xingyu Xu and Yuantao Gu. Benign overfitting of non-smooth neural networks beyond lazy training. In Francisco Ruiz, Jennifer Dy, and WnnWilleem van de Meent (eds.), Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pp. 1 11094-1110 11117. PMLR, 25-27 Apr 2023. URL https: ppppeeeiiggg mlr wrwpsprseeeeee... v20// xu2 3k. html.
Bojan Žunkovič and Enej Ilievski. Grokking phase transitions in learning local rules with gradient descent. arXiv preprint arXiv:2210. 15435, 2022.
11

---

Published as a conference paper at ICLR 2024
A
APPENDIX ORGANIZATION
A.1 Additional Notation
12
A.2 Properties of the training data
12
A.2.1 Proof of Lemma 4.1
12
A.2.2 Proof of Corollary 4.2
15
A.3 Properties of the initial weights and activation patterns
15
A.3.1 Proof of Lemma 4.3
16
A.3.2 Proof of Lemma 4.4
16
A.3.3
Proof of the Probability bound of the "Good run" event
22
A.4
Trajectory Analysis of the Neurons
22
A.4.1
Proof of Lemma A.3
23
A.4.2 Proof of Lemma
24
A.4.3
Proof of Corollary A.5
25
A.4.4
Proof of Lemma A.6
................................
26
A.4.5 Proof of Lemma A.7
28
A.5 Proof of the Main Theorem
31 ................................
A.5.1 Proof of Theorem A.8: 1-step Overfitting
31
A.5.2 Proof of Theorem 4.7: Generalization
32
A.5.3 Proof of Theorem A.13: 1-step Test Accuracy
36
A.6
Probability Lemmas
41
A.7
Experimental details
44
A.8 Additional Experiments
45
A. ADDITIONAL NOTATION
Denote the c.d.f of standard normal distribution by ((-) and the p.d.f. of standard normal distribution by d'(.). Denote (b-) = 1 2..). Denote the Bernoulli distribution which takes with probability p E 1) by (er Denote the Binomial distribution with size and probability p by [n,, p). For random variable X, denote its variance by Var(X), and its absolute third central moment by p(X).
A.2 PROPERTIES OF THE TRAINING DATA
A.2.1 PROOF OF LEMMA 4.1
Lemma 4.1 (Properties of training data). Suppose Assumptions (Al) and (A2) hold. Let the training data ((x,, v be sampled i.i.d from as Definition 2.1. With probability at least O(n -) the training data satisfy properties (BI)-(B4) defined below.
(B1) For all & E [n].
max (Ix 1e- v) < 10vlog [|]l and 1||==|1----p111 11 S 10 plon vEcenters
(B2) For each i, [n]] such that i + k, we have (()15 (ã,, Ftai 5 10 logn.
(B3) For E centers, have t nu n/41 S Ven log and (n.-n(cu+n)) S Vennlogn.
12

---

Published as a conference paper at ICLR 2024
(B4) For E centers, we have |cv + ny C-v n--ll 2 n2/2- and |n n--l 2 myne 1
Denote by Idsta the set of truining data satisfying conditions (B1)-BB4). Thus, the result can stated succinctly PXXx E (daaa Z 1 - O(nn((()))
Proof. Before proceeding with the proof, we recall that centers = { =22). We first show that (B1) holds with large probability. To this end, fix & E [n]. We have by the construction of r& in Section 2.2 that 2g ~ (Tx; Ip) for some Ts E (=111, - Let & 2% Ts. By Lemma A.17, we have
P(1E@||> > /p((+1)) I1 pl pl) S 2exp(-pl /8), 1I (0,
(A.1)
Note that for any fixed non-zero vector V/ E IP, we have Ex) ~ No, |m| Therefore, again by Lemma A.17, we have
((),, > |||/|l) S exp(- Vt 2
(A.2)
where the parameter 1 in both inequality will be chosen later. To show that the first inequality of (B1) holds w.h.p, show the complement event Fs {maxuceeners(Se,, v) [||ilB has low probability. Applying the union bound,
S(F))
Σ WE ve(
P(KE,, w) > (|pl]) 12' Union bound
s 44xpp(-85 12/2) Inequality (A.2).
Let 8 : n. Picking t /2log(16n/0 in inequality (A.2) and applying the union bound again, we have
P(U FA) < emexp(-12/2) < 6/4.
(A.3)
Next, fix t E (0. 1) and 12 2 1 arbitrary. To show that the second inequality of (B1) holds w.h.p, we first prove an intermediate step: the complement event & 1| p |||||3| pti 2/|m|||t2) has low probability. Towards this, first note that since
| = | + 1| + 2(Tx, Ex)
||a12 + |IEx + 2 Ex)
we have the alternative characterization & as
Ex {||| p ( Ex)l ptj + 2/
Next, recall the fact: if X, Y E R are random variables and a, b E R are constants, then
X + Y > a + b) S P([ > a) + P( b).
(A.4)
To see this, first note that |X + Y
<
X
+ Y by the triangle inequality. From this we deduce that + b). Now, by the union bound, have
P(IX + Y a b) P[|X
P[[ + YY|
b) < X[[
> } U {||| > 6)) < P([X > a) + P(()) 6)
which proves (A.4). Now, to upper bound PPEE)) note that
IP(E)) 1 p + 2( E > ptl + 21|m|((2)
< TP - pl pt,) PP( [Tx Ex) > to (ell)) IInequalyty (A.4)
<
exp(- 2)) Inequalities (A.1) and (A.2)
(A.5)
Inequality (A.5) is the crucial intermediate step to proving the second inequality of (B1). It will be convenient to complete the proof of the second inequality of (B1) simultaneously with that of (B2). To this end, we
prove analogous intermediate step to (B2).
Fix $1, $2 2 1 to be chosen later. Define the event Ei- := (1 2zy) (Z,, 1,y1 > $1 $1vP+ 2t ||ll for each pair i, [n]] such that i A j n. We upper bound (P(Ee)) in similar fashion in (A.5). To this end, fix i, j E [n]) such that i + j. Note that the identity (æ,, (r)) <T & + + J. + + +J
13

---

Published as a conference paper at ICLR 2024
implies that [x,, ;3)
(R., |6 Ey + + + [ /,). Now, claim that
P(E,)) P(IE] E, Zy -.l 2 s1VP + - 2f2l|a|]
P(KKT &1 > S1vP) + P(c) رزن > 42 1-]]]] + P((E) 3el > 12)|a|] - exp( + 2 exp( -p(s2 - /8) + 2 exp( (ttg/))
(A.6)
The first inequality simply follows from applying (A.4) twice. Moreover, P(I. أ,ن > tz t2ll and P(Is] 2,1 > fellall) I exp(-t3/7 follows from (A.2). To prove the claim, it remains to prove
(( f,yl /D)
I (P(( cn1 SIVP ||&/|| S /s2p) + P(1E,)11 /s2p)
law of total expectation
S <exp(->7 /2/2)) + exxp( p(82 /8).
(A.7)
To prove the inequality at (A.7), first we get P(||E)] x2p) S 2 exp( (p(ppppp222-000000000 - 1)2//) by applying (A.1) to upper bounds the second summand of the left-hand side of (A.7). For upper bounding the first summand, first let P (1(E. Ex) > 81 /D Ej) be the conditional probability conditioned on realization of Ey (while & remains random). Then by definition
((((E.: Es>1 > $1 /p ||&|||| S /82p) = Es, (1KE, & sIVP Ex)
|& & /52p]. (A.8)
For fixed Ej such that ||& S / /82p, we have by (A.2) that
((((E,,))) > $1 /1 {g) P(KEE C,3
5 (5pp-((51//11(51)))/
Continue to assume fixed E, such that ||&/|| - /52p, note that 91 v 1l 2 31 stv s2p s1/ implies
exp(-(s1 (1//ss///11,,,)))/// - exp(-(s1/
(51 /=2))
Hence, (I(EE,,))1 > $IVD1 {3) I xxp(->//222) Applying Eg, 1. ||E S /82p1 to both side of the preceding inequality, we get IP (( E,) > s1vD ||& S /22)) S exp(-s} s 2222) which upper bounds the first summand of the left-hand side of (A.7). We now choose the values for
/8 log(16n/5)) t2 222 ( /5). $1 = 2 /log(88)) (8)) and $2 = 1 /8 (olog1666 (066 /5)/p.
Recall that 8 = -= and is sufficiently large, then have
Aog( (6n / = Vlo((1622+)) 3/p S
<
by Assumptions (A1) and (A2). Combining (A.5) and (A.6) then applying the union bound, we have E((U)_1& U (ULje() s Em P(Ex) + Ex.je(): P(Et))
(A.9)
< 2n 22nexp(- + n n 12xxpl 4 + exp( = + 2 exp - 1) )] S
Moreover, plugging the above values of t1, fe and s1 into the definition of & and Eij. we see that (BI) and (B2) are satisfied since they contain the complement of the event in (A.9).
Next, show that (B3) holds with large probability. We prove the inequality involving |cs + n, n/4 portion of (B3). Proofs for the rest of the inequalities in (B3) follow analogously using the same technique below. Recall from the data generation model, for each k e [m]] Te is sampled i.dd ~ Unj, +12). Define the following indicator random variable:
(k) = {
1 if Ts otherwise,
for each & [n]. and (=)
Then we have E. 4, (k) = for each k and (k)) = n/4 for each v. Applying Hoeffding's inequality (Lemma A.18), we obtain
P PO (k)) n/4| tym)= < 2exp(-221
Applying the union bound, we have
P (mmax, Ek (k)) n/4/ t//) 1exp(_222
(A.10)
14

---

Published as a conference paper at ICLR 2024
Thus can bound the above tail probability by O()) by letting t v /log(1/8)/2, and the upper bound i S /n (og(6))) / / log(n).
Next, show that (B4) holds with large probability. We prove the inequality involving |cv nyy C-v nー-v portion of (B4). Proofs for the rest of the inequalities in (B4) follow analogously using the same technique below. Note that for each k.
1 for any 2
(k) [k)( 0:
(R)I.()) - (k))))))
yields that
(k) - - (ar()//)) (a11))//V(((()-II-o-(/ I 2.
Applying the Berry-Esseen theorem (Lemma A91 have
12
P([c + c-v - - = P( (1.) - L_((k)) > tyn) 2 24
Vn"
k=1
Let n- By (() S 1/2 (0)4, have
4
12 = -
P( Σ (k) [.)( (/1)
(A.11)
V2n
k=1
Combining (A.3), (A.9)-(A.11), we prove that conditions (B1)-(B4) hold with probability at least 1 16)( over the randomness of the training data. As a consequence of (B1), we have
p/2 S + ||xl12 - 10 10/plog(n) s ||=x|| < || + 10 p plg((n) - 2p
by Assumption (A1) and (A2).
A.2.2 PROOF COROLLARY 4.2
Corollary 4.2 (Noar-orthygonality of training data). Suppose Assumptions (AI), (A2), and Conditions (BI), (B2) Lemma 4. I all hold. Then for all < s k S n,
|cossim (I,, 2x)1 C'n²²
Proof. By Lemma 4.1, we have that under (B1) and (B2), when i + j.
|| 10 log(n)
241 -
2
|| 10 log(n) <
|xyl|
Cn²²
p
for sufficiently large p. Here the second inequality comes from Assumption (Al); and the last inequality comes from Assumption (A2).
A.3 PROPERTIES OF THE INITIAL WEIGHTS AND ACTIVATION PATTERNS
We begin with additional notations that used for the proofs of Lemmas 4.3 and 4.4. Following the notations in Xu & Gu (2023), we simplify the notation of Joos and Jneg defined in Section 4 as
Te = Joos {j E [m] >
0): Iu == Jueg = {j E [m] aj < 0}.
Active neurons/samples. We denote the set (9 of all neuroonneuron pairs of indices (i,j) [n) x [m] such that the neuron j active with respect to the sample x, at time t. In other words, we define
(+) = {ii.) j E [n] x [m] (w)))
xx) > 0}-
Next, for fixed time index 2 0, define for each sample index E [n] the subset As(00) C [m]:
(t) (j [m] (m0)) .x.) > 0).
15

---

Published as a conference paper at ICLR 2024
Likewise, for each neuron index E [m]] define the subset A,) [n]:
A = (i E [n] (wn) mi) > 0}.
Note that At(()) (resp. A neuron j) at time L. Nexi, we define the of positive (resp. negative) neurons that is active w.r.t sample i at time t:
) is the set of neurons (resp. samples) that is active w.r.t sample i (resp.
JT Jp n
Ju
IN n
(A.12)
Likewise, we define the set of clean (resp. noisy) samples that is active w.r. neuron j at time t:
n No. == Nr n nA)) for j E [m], v E centers.
(A.13)
Note that the above definitions (A.13) are equivalent to that of (4.1) from the main text. Finally, define the notational shorthand Riv == fy +
(v, -allined neurons. For each E centers and parameter & [O, }1), denote the sets of neuron indices j E [n]] of corresponding to (v, allliggnedd neurons by:
J5 = ij E [m] D > n n and m d < min c-vB - 2n+ Vn).
In other words, we have the following identity between sets of neuron indices:
JHt tj [m] neuron j is (v, -lallinndd]
Finally, we define the subsets positive (resp. negative) (v, -llaiigeed neurons:
THP = Jo n JT; Jx Tn N.
(A.14)
A.3.1 PROOF OF LEMMA 4.3
Lemma 4.3 (Properties of the random weight initialization). Suppose Assumptions (A1), (A2) and (A6) hold. All conditions below simultaneously hold with probability at least 1 O(n -) over the random initialization:
(CI) | S
3 Tmi mp.
2
(C2) [Jeos| 2 m/3 and [Treg] m/3.
The result can be stated equivalently as follows: Denote the set of W()) satisfying condition (CI) by Gw. Denote the set of (a,) satisfying condition (C2) by GA. Then IPaa GA. W0 Gw) Onn=))
Proof. Recall earlier for simplicity, we defined for simplicity J, Jroa and Jx Inag: Let 8 n Then (C1) is proved hold with probability 1-0(5) in the Lemma 4.2 of Frei al. (2022b). For (C2), since [Tpl and [Ty both follow the binomial distribution' B(m, 1/2), it suffices to show that P([Ta] m/3) 2 8. Applying Hoeffding's inequality (Lemma 18), we have
P[[TI m/3) P(JTr) m/2 < -m/6) S exp(-m/18) < 6,
where the last inequality comes from Assumption (A6).
A.3.2 PROOF LEMMA 4.4
Lemma 4.4 (Properties of the interaction between training data and initial weights). Suppose Assumptions (Al)-(A3) and (A6) hold. Given GA (defined in Lemma 4.3) and X Gáau (defined in Lemma 4. I), the followings hold with probability at least - Onn=)) over the random initialization (().:
*See (4.2) in the main text for the definition of (v, k)-aligned neurons
*See Scction A.1 for details the notation
16

---

Published as a conference paper at ICLR 2024
(DI) For all i [n], the sample x; activates large proportion of positive and negative neurons,
i.e., Hij E Jros : (w]) xi) > OH 2 m/7 and Ki E Jxag (w) (0) xzi > 031 2 m// 7 both
hold.
(D2) For all centers and E [0, ))) both {íj Joos is (v, k))-liinedd)) 2 mn 10s and (jj E "nog'' j is (V, k))-ligned) mm-10s
(D3) For all v E centers, we have {j E JPea ; Moreover, the same statement holds if "f""Ter "Too"" " is replaced with " "Jueg" everywhere.
j is (=v, - 2 (1 10n 100-200))
20 J|Prol
(D4) For all v E centers and E [0, let Tupos := n
-ū -- Trea i is (v, -aligndd.. Then 10 [v.poal- Moreover, the same statement holds if Jroo" is
Ejeröro (co - -Ty - 2 replaced with "Inag" " everywhere.
Before we proceed with the proof of Lemma 4.4, we consider the following restatements of (DI) through (D4):
(D'1) For each i E [n]. x; activates a constant fraction of neurons initially, i.e. for each E [n]] the sets J/ and Tu defined (A.12) satisfy
JJ. 2 m/7
and
[ 2 m/77
(D'2) For E centers and λ E [0, 1/2), we have min JT"nl 2 mn 10c
(D'3) For E centers, we have | U J20E 2 (1 - 10n 20 Jpl and | U T20g 2
(1 - (1-10n 20e) Jul
(D'4) For E centers and x E [0, }), we have Ejer (cv -o d 2
a Jj. where J {Jr,p J.
10
Unwinding the definitions, we note that the (D'1) 1) through (D'4) are equivalent the (D1) through (D4) of Lemma 4.4
Proof. Let ô = == Throughout this proof, we implicitly condition on the fixed {a,} } GA and {x,) E Gdats i.e., when writing probability and expectation write P(.) and E . to denote P(. } {x+3) and B[- Kay)), {z, respectively.
Proof of condition (D1): Define the following events for each i [n]:
P.
(1 m/7}; :
(1 2 m/7}.
We first show that 03 (P, (.) occurs with large probability. To this end, applying the union bound, we have
P( (P = -1P 21 Σ (P(P.) + P((T))
Note that P, and N; are defined completely analogously corresponding to when a, > 0 and aj < respectively. Thus, to prove (D1), it suffices to show that P(P.) < 5/(4n) for each or equivalently,
IP
Σ JEJ
U
m
<
8
7
2n
holds for each i e [n). where Uy := I((w!)) xx) > 0). Note that given o; and Jr. {Us}en are iid
Bernoulli random variables with mean 1/2, thus we have
P C PP( Σ (U, jEJ
U,
} < E
6 mm) I expl 2ml exp(-2ml
1
1 ) -
8
JEJ,
17

---

Published as a conference paper at ICLR 2024
where the first inequality uses [Tp m/3; the second inequality comes from Hoeffding's inequality (Lemma A.18); and the third inequality uses Assumption (A6). Now we have proved that (D1) holds with probability at least 8.
Proof of condition (D2): Without loss of generality, we only prove the results for T.p. Note that
ThP - Tup for 1i < Rig Thus we only consider the case x = 0. It suffices to show that for each
j E [m].
P(D v.j Vn) 2 8n 10c and P( min(cy,c-u) 2n vn) Sn" Sn110 14 E (+v).
10c
(A.15)
Suppose (A.15) holds for any { =22). Applying the inequality P(An B) 2 1 - P(A") P(B"), we have
P(D od d < min(cr,c-u)
2n - Vn, 1 E {+v}) 2 8n -00 20= 6n 10e
Then we have
2m 10
BI 2 6n 10s
T
where the last inequality uses [Fl |30 2 m/3. which comes from the definition of GA. Note that given {az} and { |Tupl is the summation of i.idd Bernoulli random variables. Applying Hoeffding's 's inequality (Lemma A.18), we obtain
2m 2
P(IJ.ppl S 10z S P(TT.p B[[Jupp]] S
I exp(-
S
n
n 10
n 20g
where the last inequality uses |Jb| m - |rxl S 2m/3, 20s S 0.01, and Assumption (A6). Applying the union bound, we have
P(Nveltmi.taa) [1 > m/n 10= 2 4n
Thus remains to show (A.15). Without loss of generality, we will only prove (A.15) for v = +ui which can be easily extended to other v's. Recall that X = [=I,...,
........y, xn]] is the given training data.
Let V = Xw,,, then V (0, XX Let 2 [a zn]] 2 = u//lx E [n]. Denote = Cov(Z). Then Z ~ N(O,E), E). By Corollary 4.2, we have
2
Σg 1;
|Egl S
Cn2
for 1 < s i
Denote
A1 = C+ps UN A = C-m UN+
By the definition of Gdata and (B3) in Lemma 4.1, we have
|| A < |ctps C-pul + |ntus n- (1 17) nz log(n);
(A.16)
|1 A2 C+p + n+u + C-ju + 2 2 2/ no
(nn)
(A.17)
for sufficiently large n.. Note that equivalently, we can rewrite ( as
Σ
Σ Il2 > 0).
(A.18)
REAA1
A2
Since want to give lower bound for (0) With the new expression of D. (0)
below only consider the case when |Ai [Az].
f
3 have
[A--vai
((D (0) thoj > Vn)= =
Σ
Σ
Σ
E
II
I(2> > 0)
II
I(z, s 0)
O=
C, |
B.CA
iEB,UBB
RE(A (LLAABB))L (AA)
(A.19)
18

---

Published as a conference paper at ICLR 2024
By Lemma A.1 we have
E
II iEBIU82
I(2 0).
II
I(2i - 20) 2 2
(A.20)
iE(A (.
where 7 1/2 4/(Cn). Let 7' |z| ... _ N(O, In). Denote A, EiEA, T1 > 0) CieA, 11(1 0), and na = |Ail + Agz). Then we have A B(IAI), B3|A2] 1/2),
|[A,] (1111 - 1/21)/2, and
|| Vna
2 -11+ +1++n)+ n) ns log(n)
ns log(n)
(A.21)
2/nn22 - o(n)
by (A.16) and 17). Here the last inequality comes from Assumption (A3). Combining (A. and (A.20), we have
[A
P(D +Mi > Vn) N
A
B.C. A. 1
BiA
|AA
2 1
-(2))
Σ
Σ
(A.22)
C.
|82|=k Biss+t
(27) P(A, Vn)
8
8 PP(A > Vn).
2 (1 - Cn
> v 2 (1 C
where the second equation uses the decomposition of P(A vn); the second inequality uses
|Ai] + [A2 S n; and the last inequality uses f(n) = (1 - 8/ (C is a monotonically increasing
function for
1. Note that
E[A
Vn R[A
P(A, TT) PP
V2AA22 2
V 2
/2 - 0 (22(v3+ + V=log(())) 0
where the first inequality uses Berry-Esseen Esseen theorem (Lemma A.19), and the second inequality is from (A.17) and (A.21). If (og( I v3, then ( + v log(n))) O(1/Vn) ((1). which gives a constant lower bound for P(A₁ > Vn). If / f log(n) > v3, we have
$(((V3 + +elog(())) 2 ( llg(()))
8/ 2Te 2T log(n) (op(-88log(n))
17
2 10=
8 2.T log(n)n*t
for sufficiently large n. Here the second inequality uses (xx) 2 (2xx)/(2x) for x 2 Combining both situations, we have
17
OEE
16
P(Aj
Vn)
>
(A.23)
n
10s
V
\n33
105
for sufficiently large n. Combining (A.22) and (A.23), we have
PPD Vn) (1
(0)
8
16
8
C
n 10s
n 10s
for C > 16. It remains to prove
(P(d pd
> min(C+mr, C-p - 2n VT. <
n 10 E {.
19

---

Published as a conference paper at ICLR 2024
Without loss of generality, below we prove it for μ + According to condition (B3) in Lemma 4.1, we have
min (Ctp): C-p1) 201 - vn 2 [
5n)n 6 ne log(n) vn
5 C
5
)n 2 " (A.24)
6
for C' 2 150 and sufficiently large n. Here the second inequality is from Assumption (A3). Thus it suffices to prove P(d4)) 2 n/66 S n n-10s Note that
Σ
E
1(4 > 0).
+nn
iEC+MI
iENti
Denote
A
Σ rEC+.1
1(2) 0). 1(2 (14>0) 0)
Σ
iEN-
Following the same proof procedure for the anti-concentrrtiio result of +m we have
(0) +in
m 5
<
(2x2)ctm +D+1 P(A) 2
n ))
where 122 11/2- + 4/(Cn). According to condition (B3) in Lemma 4.1, we have c+m - n+ps S (1/4 - 2n)n + /ne log(n). It yields that
BJAJ] C+m n
S (1/8 7)m + Vne log(n) S n/77
2
Applying Hoeffding's 's inequality (Lemma A.18), we have
P(A) 2 n/6) < P(A2 - E[AA] 2 n/42) S exp(-9(n)).
Combining the inequalities above, we have
8 yoim C0IntRIPPPP P(A) 2 n/6) = exp(-I((n)) S
P( 2 n/6) s (1+ Cn
(A.25)
n 10 .
where the equation uses (1 + ( th4n1 & (1 + 8/(Cn))² v exp(8/C). Now we have completed the proof for (D2).
Proof of condition (D3): Without loss of generality, we only prove the results for J20g U J20g By Berry-Essen theorem, we have
P(|4,1 S n 1/2 202) P( A ]
F[A,]
2
E
E
2
+
VTa/2
/
n 20
VTAA2
/2
20
2
1
< 2 n0 ) [(0) + +0(
< 4n - 20c
Vn
where the first inequality uses (b) - ((a) < 2(4((bb-)))) aa/2) - (0)), 2 a; the second inequality
uses & (xx (00) I (0(O)r, x 2 and 20g < 1/2 yields that
P(D (0) | S n 1/2-20 < 2T,(A, S n 1/2-20c) < 8n 20s
where the first inequality is from Lemma A.15. Combined with (A.24) and (A.25), we have
P(ID > n 1/2-20c 20c
mi. c - 2n - Vn, {+11))
dz
> P( P > ni 1/2-20s 20s
d0) < n/66 E (+
> 1 - 8n 20e - 2 (P(n) > 1 - 9n -20c
9n
where the second inequality uses (00 (0)
(0)
-i and PP A.) = 1 = 49) > 1 -
Σn P(A)). Note that given {az} and {x IJ Ju.u is the summation of Bernoulli random variables with expectation larger than 1 - 9n 20g Applying Hoeffding's inequality (Lemma A.18),
20

---

Published as a conference paper at ICLR 2024
obtain
P(( U - [Tb
10n -200)
$ P
20c U
-
J
[BI 20r
P U J 20c
-npippl < -IJ 20c)
10e)
- n
Jp [1 9n 20c) and the last inequality is where the first inequality uses ] U JJ-mpPP from Assumption (A6) and 40c < 0.01.
Proof of condition (D4): Lastly show that (D4) also holds with probability least 1- Onn=)) Without loss of generality, we only prove it for J Referring back to the definition of J in equation (A.14), it is crucial to note that it solely imposes upper bounds on d Consequently,
the average of (0)) dimi- in T is no more than the average of d in Jr, which imposes no constraints on d0)) probability "1s"" Armed with this understanding, when [ > 0, we have that with
(c+pr n4#1 dmnis) 2
(0))
(ctm
n+tp1 - d
Jf
Jjp
JEJF
JEJ
Thus it suffices to show that
(ctp)
(0)
n 2
(A.26)
Jo
10
JEJ
with probability at least 1 - 06)). Note that given the training data X, {d )
are i.idd random variables with _(d m = (c-p1 n/=1)/2 which comes from the symmetry of the distribution of
1m
w) Then we have
1 E|c+m - 14 d = C+p - n- (c-pi - nn-,1//22 (g -5)) - 5v ne log(n) N
n
9
(A.27)
Here the first inequality uses (B3) in Lemma 4.1 and the second inequality uses Assumption (A3). Applying Hoeffding's 's inequality (Lemma A. A.18), we obtain
n 11tm1 - d <
(c-m
JTp
10
JEJ
Σ
(a(0)
-Mn Eld =1 > (c+m) n+#
E(d
pp ()) [[Tp]]
10
jE.
n [P]Tpl
(0)
(l
#1 ,])
[Fl]
S
exp
S 8,
90
4050(c-m) + n2 nmp1)p
JEJ
where the first inequality uses (A.27), the second inequality uses Hoeffding's inequality (Lemma A.18) and the bounds of d (0) (A6). It proves (A.26).
"".3"" i.c. -n- I (00) d=p C-ai and the last inequality uses Assumption
Remark A.1. the proof of (D2), note that when E In. 2; are independent with each other. Then (A.15) can be proved by applying Hoeffding's inequality (Lemma A.18). In our setting, Σ is close to the identity matrix, which means that {2]] are weakly dependent and inspires us to prove similar results.
21

---

Published as a conference paper at ICLR 2024
A.3.3 PROOF OF THE PROBABILITY BOUND OF THE "GOOD RUN" EVENT
Combining the probability lower bound parts of Lemma 4.1,4.3 and 4.4, we have
P((a, w(, gguen)
P(a GA, X Gdatn; (D1)-(D4) are satisfied) F(W(O) Gw) 2 P((DI)-(D4) are satisfied a E GA. X E 9amu) E GA. X E Gdwa) O(n On =))
21 (n) On
(n = O(n=)),
as desired.
A.4 TRAJECTORY ANALYSIS OF THE NEURONS
Let 2 0 be an arbitrary step. Denote (1) 2; decompose (2.2) as
=- yis( W(O), and 11 (1) -1/2. Then we can
w ((++1) w/)
001
Ei=1 ('(wW)) ,.,))aa; aai E h10 d(O((( (x:)))) ,.))),... 2n
(A.28)
Remark A.2. When 1 (1) is sufficiently small, we can use 1/2 as an approximation for the negative derivative of the logistic loss by first-order Taylor's expansion and we will show that the training dynamics is nearly the same in the first O(p) steps.
Lemma A.3. Suppose that Assumptions (AI)-(A6) hold. Under good run, for 0 t 1//1/p/p))))
2. we have max | < 2/n3/2
Lemma A4. Suppose that Assumptions (A1)-(A6) hold, Under good run. for 0 S $ 1/////p/pp
2. we have that for each k E [n].
(wo
wy xk)
0gg
[u ((10)) (0) .kk)pp + Va, 1 P2.j [//1
2n
4a
(6' ((wg)) ; Cn
||11
and
(A.29)
n//2 Vm
30
(10/tt1) wij
v)
aa;
yo ((th |||0||
S
5a
||1|1
(A.30)
2n
3/2
m
where Cn := 100ggg)), I& E centers is defined as the cluster mean for sample (Ik, yk), and yo is defined as the clean label for cluster centered at v (i.e. yv = 1 for v E { Ve = 1 for E (+22
Taking a closer look at (A.29), we see that if a,ux > 0, and øy activates neuron wy at time s, then Tk will activate neuron u/ (1) for any E [s, 1/(Tripx) 2). Moreover, if ayylk 0, and zy activates neuron Wj at time 8, then Ie will not activate neuron wy at time s + 1, which implies that there an
upper bound for the inner product (w])) .rn). These observations are stated as the corollary below:
Corollary A.S. Suppose that Assumptions (AI)-(A6) hold. Under a good run, for any pair (j k) [m] x [n], the following is true:
(El) When ayyk > 0, if there exists some 0 < s < 1/(/npa)) - 2 such that (w) .K) > 0, then for any s < < 1/(/rpp)) - 2, we have (w .nn > 0.
(E2) When d,yk < 0, for any 0 L 1/(vmpp)) - 2, we have that (w)g) xxa) S vm 11
(E3) When ajykk 0, for any < t S 1/(V/pa)) 3, we have that (w0) .kk) 0 implies (mu (1+1) .xk 0.
22

---

Published as a conference paper at ICLR 2024
A.4.1 PROOF OF LEMMA A.3
Lemma A.3. Suppose that Assumptions (A1)-A66 hold. Under good run, for S 1///1////// 1/(vip)))
2. we have maxig[s] | - 2/n2//2
Proof. It suffices to show that for 0 <t 22
max
11
1
2ap
(t + 2).
sE[m]
n
We prove the result by an induction on l. Denote
200
P(t)
max Ih
(7)
(t + 2),
VT VT St.
iEin|
When = 0, we have
|n | S
pwinitt V 33m
v /3allal1
4ap
2
4nm
by Lemma A.10, Assumption (A2) and (A5).
Thus P(0) holds. Suppose P(t) holds and t <
1/( (Vppa)- 3, then we have
16
2ap
(T 2) S
2
2
2 Vn
Vr St t,
vn
2
which yields that maxie[n] 93 (r)
< 1. Further we have that for each pair (j, k) E [m]
|n].
(7) (wwftt)) - -0,² xxn1
(r) 0' ((w)) (T)
9
1=1
nym
a
max 1E
(2p + 2m|| <
nym
4ap
where the first inequality uses ||=;11² < 2p, I,)ر < 2μ²², which comes from Lemma 4.1, and the second inequality uses Assumption (A2). yields that for each pair (j,k) [m] x [n]]
(w (t+1) mami S
(T) xx)1+ Σ Kun (uob+tD>ww7
.ma)l S nym
4ap (1+1)+
2p 10
4ap (I+2),
nym m
where the last inequality uses Lemma 4.3 and Assumption (A5). Then we have that for each E [n].
[sa W(+1))) ((Tt:)()))))) s
m
a (w (t+1)
xxحi S vm max
(w ((+1)
xx)) S
4ap
n (1
+ 2).
jE(m)
By [1p(1+) 1/21 < |=|/2, Vz, we have for each E [n].
16 11 $ マタ 12 (0
2 f(z;; (((+))] v 2ap (t + 2).
Thus P(t + 1) is proved.
As consequence of Lemma A.3, have g;
(1)
[1/4,1] for 0 1//51 -2
23

---

Published as a conference paper at ICLR 2024
A.4.2 PROOF OF LEMMA A.4
Lemma A.4. Suppose that Assumptions (A1)-(A6) hold. Under good run, for 0 S 1 S 1//////////// (V/pp))
2, we huve that for each & E [n],
(wf (t)
m xb)
aa
[uko' ((u))) (ouf),xx)))p + Vas D 2nd [lll
2n
4a
oo' [d'((ug)) xx)))P + C, 1.99 Co ||11 },
and
(A.29)
n2 mm
30
w/
v)
aaj
yo D |1/13
S
5a
||w|12
(A.30)
2n
n3/2 m
where On := 10 log(n), Ia E centers is defined as the cluster mean for sample (Ih, Wa), and y defined as the clean label for cluster centered at v (i.e. Yv = 1 for v E ). yu = - for E (=
Proof.
First we have
(aaj n
h (1) Φ' ((w)))
((uff,,;)))lee as)) mx)
20
Σ Φ (1u (1) (2)))
1//22 m
2a
<
Oo' (1w) .xx))<>>> 11? + > ist
(az xx)] (A.31)
n/2 mm
4a
(o' (11u) (0)
zx))p n
5/2 mm
where the first inequality uses max; 11 < 2n 3/2 which is from Lemma A.3; the third inequality uses |arx]1 12= S 2p, Ka,, x.ح] S 2 |ul2 which induced by Lemma 4.1. Next we have the following
decomposition:
Σ =m mx) (|=xa11 - -- ||a|13) + Σ 6((10)) ( (e:))e(( (B,, ak))
o' 10, ,,))) ((,,,,,))
is
+ + m (1) ((w)), [x))( + 11x113) + Σ
& ((uug)) xx)) (e (T,
(A.32)
*f
=m 11) 01 - p 1111112) ) + Σ
8' ((uu))
(Z., Tx))
is
yn ( (mwg) (1) rx))pp + yz, 2 D
||a112
Σ
0'(1u))) (1) xI)) Eb).
where the second equation uses the definition of (10 Recall that C, 10v/oo((n) Combining with results in Lemma 4.1, (A.32) yields that
Σ d' ((wW)) ar)) (x)))))) sn) yyx o' ((w))) ax))) [ S &nCn 2nC,, VD.
(A.33)
where the first inequality uses (B1) and (B2) Lemma 4. 1 and the second inequality uses Assumption (A2). Recall the decomposition (A.28) of the gradient descent update, we have
(w) (1+1) -w/ xx)
aaj
o' ((ww))
Ix) + aaj
(4) ((((w)) (1)
(1))), xk)
2n
(A.34)
24

---

Published as a conference paper at ICLR 2024
Then combining (A.31), (A.33), and (A.34), we have (1)
(u0) (1+1)_
adj
( (1) )x + yea D [
wj rk)
2n
4a
[d' [d'((1]]) ,xx))p
+ (((1u)),xx))p
+ aCrvD
n5/2 /m
vm
4a
od' [d((w)) (0) .xx) p+ n||la12 Cnn
||,112
12//2 m
4C
4a
Od' ((w) (0 mx))p 2 |||l?
n// vm
3C
where the second inequality uses Assumption (A1) and the last inequality holds for large enough n.
Now we turn to prove (A.30). Similar to (A.34), we have decomposition for (ur} (ue((+1) w v):
(w) - w 1)
aaj
o ((w) (u))), +
aaj
h d'((w) (0) ,x:))
2n
Similar to (A.31), we have
aaj
h
(1)
(M) (,,,, (x))) ((w)
4a
||1l
n2/2
n
by Lemma A.3 and (a,, 1 S 2 2|| which induced by (B1) in Lemma 4.1. Similar (A.33), we
have
((w m.) (u))), - nn D.D |||l
Σ &' ((w) ,a)))e -2,,),, 2,.0) S nC ||
(A.35)
by (B1) in Lemma 4.1. Combining the inequalities above, we have
(mn(t1) - y v)
ada yo
4a
|| +
aCn,
|| -
n2322 m 1/1212
5xx
2n
n3/2 \m
2/m
for large enough n. Here the last inequality uses
||d|| 2 0.51 VD 33 1.51 ||51||1
which comes from Assumptions (A1)-(A2).
A.4.3 PROOF OF COROLLARY A.5
Corollary A.S. Suppose that Assumptions (AI)-(A6) hold. Under a good run, for any pair (j. R)) E
[m] x [n], the following is true:
(El) When ayyk > 0, if there exists some 0 S < </<<<///ppp (Vppp) -
such that (w))) .I)) > 0, then
for any < t < 1/(/ppa)) 2, we have (w (1) .xn) > 0.
(E2) When a,yh < 0, for any S 1/(/1/pppp 2, we have that (w))) ,Ia) S 7m ||
a
S t 5 1/(Vnpa) - 3, we have that (w))) xx) > 0 implies (E3) When a;8b < 0, for any (w) (t+1) xz) <
Proof. (E1): It suffices to show the result holds for = s + 1, then by induction we can prove it for
all < t < 1/( 1/(vrpa)) 2. Note that a,yk 1/ / m and (au) ,II) O, by (A.29), we have
(ww (0+1) - -u xx) 2 2m
a (p- (
4a
[p+
Cnn 1.99 ||1|
>
op
> 0, (A.36)
n5/2 mm
3C
4n 4n
25

---

Published as a conference paper at ICLR 2024
where the second inequality uses Assumption (A2).
(E2): We prove (E2) by induction. Denote
Q(t):
(w)) (1) Tx) I ||m1
When t = 0, by the definition of good run, we have
kw] (0) mmmi S I || |I < [[W(()) |F /2p S WiniP/33 S
Cnvmm |aal
a
(A.37)
where the second inequality uses Lemma 4.1; the third inequality uses Lemma 4.3; and the last inequality is from Assumption (A5). Thus Q(0) holds. Suppose Q(t) holds and 1/(1 If (w)) .za) < O, we have
(1+1) .x) S (m)) [1+1] 10))xx1)) (w)tt))
xxx) S
a0jUaa (et)
||lll2
4aCn
||a11
S
a V m
||1
2n
3.n0.51
vm
where the second inequality uses (A.29) and o' (1((u) (1) xn)) = 0; and the third inequality uses D S and n is large enough. If (ur vR) > 0, we have
(uu) (1+1) w/ xnd S 2n m
a
(p --]1|111)))) + n/ m
4a
[p +
C |
30
8ap
(p n||ull?)
2n/m
n 5/2 mm
where the first inequality uses (A.29) and o' (1w, xz)) 1; and the second inequality uses Assumption (A2). Combined with the inductive hypothesis, we have
(w)
(1+1) (t) (++1) .xx) (wo (t) .x)) + (wj wj .xk) S
Vm ||
2nym (p- m||
8ap
n/ m
by Assumption (A2). Thus +1) holds. And (E3) is also proved by the last inequality.
A.4.4 PROOF OF LEMMA A.6
Since the analysis on one cluster can be similarly replicated on other clusters, below we will focus on analyzing the cluster centered +/4 Given the training set, D(u) T is function of the random
initialization u D(u) tri. plays important role in determining the direction that w/)) , 1
aligns with and the sign of the inner product (w ,Kk). For 3g E (#:4 ya 1. Then for each
t < 1/(/mpa) - 2, (A.29) is simplified to
(auftt11) - -w)",, ,xx) aa5ykP
4ap
+
|| when (w)))
(,.) (,x)) 0: (A.38)
2n
5//2 mm
2
(wu(tt))= -w) xxn)
(0)
onj
((0)
<
4aC,
when (wf"), xx) 1 0. (A.39)
2n
3Cn
Vmn
Here Cn = 10 [og(n) is defined in Lemma A.4. We will elaborate on the outcomes for neurons with a; > 0 and aj separately in the following lemmas.
Lemma A.6. Suppose that Assumptions (AI)-(A6 hold. Under a good run, we have that for any j E J20g (or equivalently, for any neuron j E Joee that is (μ1, 20e)-alignae the followings hold for 1 < < 1/(npp))
(FI)
c4 = C+ai
ci = c(0)
c(t)
NOO
- = Ø; D C+ps - -n-pi ())
26

---

Published as a conference paper at ICLR 2024
(F2)
(w]) - wj 11 p1) 2 Anym +uي ||
Proof. Given j E JTps -pi.P. when & = 0, for z E c Ctiii we have ajy' > 0. Thus by Corollary A.5, we have
F&
c +pa.g
0 0 1st$1/(/npol 2.
(A.40)
Similarly we have that for o% E c0 Hi
Tk E ((t) C-p
0 1/151/(VTpp - 22
(A.41)
EN10) -m & N - since ajyk < 0. and for ay E
Next for TY E C+mi /c + we have
(wn)) (1) -wg
(0)
xkm 2
aa D + |||12
4uCn
|//112
2n
3Cn0.01
vmn
(A.42)
2
a
Vmn ||a112
4aCn
/mm || 2 2n /mn
2n
30n0.01
where the first inequality is from (A.39); the second inequality uses D +uis > n 2/222-202 20e which is
from j E J7th pi and the last inequality uses 40s < 0.01. It yields that
(w)) (1) mx) 2 (w) (1) w .4) || | || 2
4n20c 20 vmn ||a112
a
1lul/ > 0, (A.43)
where the second inequality uses (A.37). Thus we have
C+ps |c
C
Combined with (A.40), obtain (()
+ C+pi. Then by Corollary A.5, we have
c(t)) Han C+prs 0 0sts 1/(Vrpa) 2.
For 2 E (C-p, ( U (N 2i W Following similar analysis of (A.43), we have
(w))) xx) < (w)) -w
w
(0)
Ix) ( (0) ||-1|==1| - 2n
-A Vmn
||||
|| (2) < 0. (A.44)
Cn
Thus we have C-p, (c & c()) G and N-, With NV(1 - Combined with (A.41) and A & A(l)) - we obtain
c (1)
Mi. - A
c(0)
It yields that
( FPn C+p1 - (1 tangl 1 > C+ml - n+ 0 >
Vn,
where the last inequality uses d0)) d < minfc+ C-ui 2n
Vni
and
Vn (0)
d
Vn.
C+p1 7
Thus (FI) holds for t = 1. Then (F1) is proved by replicating the same analysis and employing induction.
For the inner product with the cluster mean +P1, by (A.30) we have
(w) (w] (1+1)
wi (1) u1) 2 2m f slll1 33/2
a D
50,a
||al 2 4nymm 14
Q
D
27

---

Published as a conference paper at ICLR 2024
where the last inequality uses D4p O.
A.4.5 PROOF OF LEMMA A.7
Lemma A.7. Suppose that Assumptions (A1)-(A6) hold. Under good run, for any j T n J20r (or equivalently, for any neuron j E Jnag that is (#11) ( -aligndd,,,, the followings hold for 2 t 1/V/ppa)) 2.
N1t +HI.f
-N
N -1 -p1.5 = -Nipii
N
(A.45)
-n - -m (t - 2) S Σ
D vj In+ Am 1 2), VE {+/1).
(A.46)
8=0
where Ap, = |n+ps n- + vn.
Proof. For a given v E {#ui) suppose j E Jea Then we have
0j <
D. n 1/2-20e ( S (in, - 2n - vn}
(A.47)
according to the definition (A.14). Note that we study the same data as in Lemma A.6 and only sgu(a,) is flipped in the trajectory analysis compared to the setting in Lemma A.6, our analysis in
the first two iterations follows similar procedures in Lemma A.6. For 2x E c2 _ by Corollary A.5, we have
Ozlk < 0,
(w)) .K) < O.
(A.48)
For z E NN UN( - azyk > 0, by Corollary A.5, we have
(w))) (8) , > 0
(A.49)
for any & S 11(Vnpoo)) - 2. For Ja E (c,, (c)))) (Nr (WW))))) similar to (A.42), we have
(w)) - w xk) S - ( aaaj (0)
2n tam |llll
3.100011 /mn I| S
4aCn
a 4n
|| < 0,
/mn
then similar to (A.43), we have
(wf) (1) mx) -w
(1) m. (U)
[.x)++11u ,,,,,,,,,,,,,,,,,010 S
||a11
Cn/m 111111 <.0 (A.50)
20s /mn
For ok E (C_, (C-u(c_o),) U (N ( similar to (A.44), we have
100,,2) 2 (w))) w0 xx) |||| || 2
a
||
|||112 C 0. (A.51)
a Cr
2n Vmn
Combining (A.48)-( have
2 Ø; co) = CLviclo
Nn = -N0, M = -
(A.52)
Thus by the definition of (()) have
D æj = - - c-v + 1c + +n S -1N10 c-v + d0)) + 2n_
(A.53)
It further yields that
D + - ccp + 2n_,
-00) < -c_ + 2n_, , + d00 < <-vn.
where the first inequality uses (A.53) and the definition of D. and the third inequality uses (A.47).
28

---

Published as a conference paper at ICLR 2024
After the second iteration, for x E No We (wr (0)
.n) ,n)) < 0, (wn)) ,k1)) < 0. Then have
(w)² wj (0) .R) 2
a
1 U.j + D |aal? 3Cn0.01 vmm |111
4aCn
2nym m
4aCn ||a1
2mmm
3Cn0.01 vmn
where the first inequality uses (A.39), and the second inequality uses D + DDd < VT. It further yields that
(2) (un rk) 2 (10) - (0) mk) , I 1111/x211
2/mm |
4aCn 3Cn0.01 Vmn || Cn/m || > 0.
(A.54)
For 26 E N20 UN note that ajyk 0. Then by Conollary A.5, we have (2) mx) > 0.
Combined with (A.54), obtain No3) No. NNv. Again by Corollary A.5, we have that for 2 S 1/(pppp)
No = No, N = N-v,
(A.55)
i.c. for 2 neurons with j J2h U J are active for all noisy points in Wips which proves (A.45).
For I& c( -- note that ajyk 0 and (w, xn)
O. Then by Corollary A.5, we have (w{3),z .Ix) O. For 2b C-oic by (A.52) we have (w,)) Ix) > 0, (w xxk < 0. It
yields that
(2) (un
W xK) <
(0)
2n/m (p+ [ Dul |||||| |2)+
4ap
||+|+
4aCn 3.n0.03 Vmn
|/111 S
op
5/2 vm
2/m
4nym
where the first inequality uses (A.38) and (A.39), and the second inequality uses Assumption (A2). It further yields that
(w)) xL) {w)) - w (0) rk) + +1w0g 10 || || S
op Anym
Cnym
11u| <0
(A.56)
by Assumption (A2). Thus we have c!3) 83
For Ik c. (w,) mR) > 0, (w .kk < 0, which is similar to the setting of C_, e_v\e -v.3" Repeating the analysis above, we have
(1)
(c(1)
(2) xk) < O. (w)
For Th E C. Ac.. note that (w))) (0)
xn) < 0, (1)) (1) xn) < 0, then we have
(w) (2)
wj xK) 2
(0)
a
(D u.j + D
( ||lll
4aC,
2n m
3Cn0.01 Vmn
a
||r||
4aCn
|||| > 0,
2 2/m
30n0.01 Vmm
where the first inequality uses (A.39) and the second inequality uses (A.53). Combining the inequali- ties above, we obtain
c(2) = C, = ( V
c =0; Ø; N19 No: N -N
(A.57)
Combining (A.52) and (A.57), have
Σ
D Vjj
Cy C_v ny + 3n 2|,
29

---

Published as a conference paper at ICLR 2024
and it yields that
Cv C-v 3n, 3n S ΣD S Co C-u + +3n ny-
It remains prove (A.46). It suffices to prove
co vo-2c-u-Ano -n 1 (1- S s=0 0
D.d S Cooth Nv) + (1-2), (tm))
since 2c,, C-p 4n ny I and - 2c_v 4n, + 3n 2 -n by Lemma 4.1. Without
loss of generality, below we only show the proof of the right-hand side. Denote T {t E [T]. 2 3. (D > } [t, 8 < t2 < ... tg. To prove the right-hand side of (A.46), it suffices to show that the followings hold
S c2 + n + A (8 14);
(A.58)
4+-1
Σ
D S 4 μn (1+41 4)
(A.59)
t=
for any i E [K] and all s [1, 2]. (A.58) directly follows from the definition of the set T and the fact that D S + n for any j. t. For a given t,, tq E T. we have ((ta) (A.39), we have that for any 2 E (x\cc! (j),
vg > App 2 n. By
(w]) (utt xK)
S ,xx)
wj (ts) ..kk) S
t(hh) 2m m us
a
|111
+
4aCn
vmn ||
3Cn
a D 4m vi ||u|| O,
(A.60)
which implies that w (1:+1) still inactive for those zg that didn't activate w (4)) For any Te E c(ta) since 0,U2 < O. by Corollary A.5, we have
al|A|1?
(w)
(1)) and S
vm
Combined with (A.38), we have
(un (nngtt3) rkk) = (10) (1(t1)
(1+
(1a) (xk) + {wu) xzk
wr
S op
4ap
+ 2/m
3c
|11112
S
op
0
(A.61)
2
n//2/mm
4n 4n
where the second inequality uses Assumption (A2). Combining (A.60) and (A.61), we have CD Ø, and
(t++1)
(1+ (w) (w .R)
<
a 2m
D. ||ull2 1ullllul12
4aCn
||l
(A.62)
3.01 Vmm
for all D& E Cy. It yields that D(+++1) vj = (c(t++t) (c(t)tt)) + nh - ny = -1c [cs + n - n2 S 1nmm - -1-p2l
where the first equation uses (A.45). It implies that 4,+1 - 4, > 1. Let t7 = minft E N tq + < t S 44+1, c!o (i) # Ø). We claim that E well-defined for each because ch.s (i) # Z. Otherwise we have D < intt - < Am. which contradicts to the definition of the set
T. Thus to always exists. Choose one point from the set c/47) and denote it xτ. Note that for any
30

---

Published as a conference paper at ICLR 2024
Le [4 1,1g 11. we have c) (j) Ø, D
|n+11 -n-11b and by (A.39),
(w) ( wj 21t
a
|ull
3Cn0.01 /mn 1
4aCn
2n 2n/m m
Combined with (A.62), it yields that
4-1
0 S (my (f:) xx =
Σ
(w/(t1))
10) Tkk + (10 ( ,
a11111
4
isj
D
3C1 (tg 1:))
2nmm
t=
It further yields that
4/TCC
Σ
D Mj
S
301 (tz -4) & vn(t, t.).
t=
If 6 = f141, then we've 'e proved (A.59). If t < tit1, then we have
4-1-1
6-
f4111 -1
Σ
D
(0)
D S Vn(t" 44) + (1+t1 10) S Au (1i+1 - -44),
t=
which proves the right side. For the left side, similarly we denote T_ {t E [7], 3, D
- [t t to
tk. Following the same analysis, we can prove that the
followings hold
D vj Σ
c- ny An 4))
Σ 0m (2)
for any i [K] and all E [t,4441 - 2]. It proves the left-hand side of (A.46).
A.5 PROOF OF THE MAIN THEOREM
We rigorously prove Theorem 3.1 in this section. The upper bound of in the theorems below is 1/( 1/(/npa)- 2, which by Assumption (A4), is larger than vnn the upper bound in Theorem 3.1.
A.5.1 PROOF OF THEOREM A.8: 1-STEP OVERFITTING
Theorem A.8. Suppose that Assumptions (A1)-(66 hold. Under good run, the classifier
sgn(f(x, W(1)) can correctly classify all training datapoints for 1 < t $ 1/((npa))
Proof. Without loss of generality, we only consider datapoints in the cluster C+p, UN+ According to (D1) in Lemma 4.4, have that under good run, (T) 2 m/7, (0) 2 m/7 for each i E [n). For Ia E C+ by Corollary A.5, we have
(w0))) xx) >0
for all j Jh and 0 S 5 1/(vnpa) - 2: and
(w/)))
(4) , xx) <
α |
31

---

Published as a conference paper at ICLR 2024
for all j Iu and 0 S s 1/t11//////)) 2. Then for 1 1st 1/ 1/(vpppa) - 2, we have
Σ ==
(0 xxx) 2
Σ
((w)
(1) xx)) - Σ
vm (l( /
1
snm)
Vm
jEJ, (0)
5:a, <0
[1) (s+1) - xkm - Σ
a |1a|
m
J:ay <0 m
jEJ, (,
s=0 (0)
apt
(h)(
Fu ||
4nm
apl
0: /||1/|1 > 0.
28n
where the first inequality uses Φ(x) 2 0, Vr; the second inequality uses the definition of J (2((0)) and (E2) in Corollary A.5; the third inequality uses (A.36) in Corollary A.5; and the last inequality is from Assumption (A2). For EN+ similarly we have
Σ ( (8) _0 1=1
Σ
Mau xnn) +
Σ
1
(1) ,xx))
Vm
Vm
je,,5.10)
3:a;200
S
m (w))
w
Ikk +
vm ||112
jEJ,
(0))
j:ny>0
apt
S
a aa|m||1 ?) < O.
28n
Thus our classifier can correctly classify all training datapoints for 1 < t < 1/(yripa) - 2.
A.5.2 PROOF OF THEOREM 4.7: GENERALIZATION
Before proceeding with the proof of Theorem 4.7, we first state technical lemma:
Lemma A.9. Suppose that [[]] > 0. Then there exists constant e O Ouch that
(-c
E (BenNord [ax:W]]]]] W)] )
((r.i)~ 1 sgn(J ( W)))
max 2 exp #'centers
||
Proof. It suffices to prove that for each E centers,
EE W)] ))
(E W) < 0) S 2exp
(-c
(A.63)
W
Then applying the law of total expectation, we have
P(t.ū) (T + sgn(S W)))
Σ
PeNN(u.L)) #
VCcenters
EEnN W)]
2
exp
W
vEcenters
BewN( W)]
2 eex
vEcenters
|
Since for each v, NN, Ip)) is 1-strongly log-concave, plug in À = 1 in the proof of Lemma 4.1 in Frei et al. (2022b). Then (A.63) is obtained.
32

---

Published as a conference paper at ICLR 2024
Our next theorem shows that the generalization risk is small for large t. Recall the definition of J1 and J2. we equivalently write them as
JI Jta. 1i E [m]] : 0g 0, (1) tMisj 122-2200
d -mi
C-m -2 vni;
J2 Jtn, U JJ (i [m] 0j < 0, D (0) n 2222-200 20c
min (an((cr)) (C2, c-u) 2n Vn, {+m)))-
TTpe N: and J20: -iNN Here J P run,
are defined in (A.14). By Lemma 4.4, we know that under good
10
JIIZ
JJ 2 (1
20e Jk
(A.64)
20
Theorem 4.7. Suppose that Assumptions (A1)-(6) hold. Under a good run, for On 10e < 1 S vn. the generalization error of classifier (f W())) has an upper bound
P(z. (y # sgn(f (gu(((x;W()))) z exp (-2 ( nn= 1-1-00e [|
p
Proof. Without loss generality, we consider æ follows ((+11.pp)) Then we have
m
Ex (v.
W(th)]
Σ
Ey Bo
(u x>)]]
Σ
((w)) [[x]) - E Ex (o x))
(A.65)
m
jiay 20
5:a=<0
(w pn)
Σ
( ,z)],,
jUE
j=ay<0
where the first inequality uses Jensen's inequality. By Lemma A.6, we have that for j E J1.
, .11)
Σ 1 (0+1)
w .11 (10) .1)
1-1
((s) |lll? Wini v 3mp//2 ||
4n m
(A.66)
==0
al a||14|?
1/2 20k + (C+u) - n+t1 - (0) d 1)] Winir v 33mpp221llll
4m/m
a|l
(C+p. ntp - dt - 1))> 0,
4m
where the first inequality is from Lemma A.6 and (C1) in Lemma 4.3; the second inequality uses the property that for j E Tr. D(e) D 2 C+r, n4 and the third inequality uses Assumption (A5), It yields that
(0)) (j), 2 1, which is also from Lemma A.6:
Σ (lwof) .1)) 2 allpI|2(t 1) J:EJ
d (1) - n+ 2 allx - 1)
Σ JEJ
(c+m)
40 [J1,
4m
(A.67)
where the last inequality uses (D4) in Lemma 4.4. For the second term in (A.65), note that we have ((Ax) = Ao(x), VY > 0, and by Jensen's inequality, Φ(œ1 + x2) S Φ(z1) + d(I2), VIr, T2 E IR.
Then we have
Ex [o((w, x))] (((w,))) + Ey Bo((w p1)) = (w, (1) +1 2T |0
(A.68)
33

---

Published as a conference paper at ICLR 2024
where the last equation uses the expectation of half-normal distribution. By Lemma A.3, we have 9i S 1, and
|mw] (1+1) w/)')
aaj
(1)
o' d'(ww xz) yarl
(8)
91
max Di nym xE(m)
(8) | 12
(1i,
20vp Vmn
05151/(V7pa) 2,
where the last incquality uses ||=;112 < 2p. Assumption (A2). It yields that for each E [m].
Kar, xzحهه I 24², which comes from Lemma 4.1, and
|
Buy (rafrt)) wy
(
111
(0)
2tt/t Vnm
+ ||
30 Vmn
(A.69)
Then consider the decomposition of where the last inequality uses Lemma 4.3.
Eginyeo ((w))) .i))::
Σ M(w)
10 .1))
Σ
p((u))
(0) .1)))+
Σ
MMu, .1)).
33,<0
JEJz
JEJsJ4J2
For the first term, we have
Σ
(lu, m1i S Σ (ug
(0)
(0) Miil
JEJa
JEJ2
(w) (w)ett) (1+ w ,12i + +Kuu |t)l|
x
jE
5allul1?)
((o)
nymnn )
+ wint v /3mp/21l//
(A.70)
mm/m
JEJ
(n + A us (t (() + 1011
Σ
+ Wint / /3mp/21lall]
2m m
n
JEJ2
alll!1
a||A|1? [n ++ + 1)(t 2)][52),
Σ JEZ2
in +1+ (A,, + +)( 2)] <
2nmm
2nym
where the third inequality uses (A.30) in Lemma A.4; the fourth inequality uses Lemma A.7; and the fiveth inequality uses Assumptions (A1) and (A5). For the second term, we have
Σ
o((u)) (c) (((u))pp))) .1))
JEJn.EJJ
Σ
Σ (un ((uf(tt))
w." mmmi + Kit (0) |nni]]
jEJn.jtJa
s=0
2|| (
50a1|al1 ) + wini v 3mmp21l/u])
2mmm 1
(A.71)
nymmn
JEInitJJ ==0
at(max)e(#M)) fcr + + + 1)1
n/m
JEJn.EJZz
10at||/412
m/m (IJrl - (J2)) <
20 m [Tll.
34

---

Published as a conference paper at ICLR 2024
where the second inequality uses (A.30) in Lemma A.4; the third inequality uses Assumption (A5) and [D(]] I max n--,, C-v + nw), which comes from the definition of D the fourth zij inequality uses + Th- + for all centers, and the last inequality uses (A.64). Combining (A.68), (A.69), (A.70), and (A.71), 1), we have
Σ
( (0)
x))] x v
Σ
(0) (1)) o((ug)
Σ
2
juacc0
faay<D
Σ
(t)
E
al(uu} (t)
.2)))
Σ
(4)
2T
JEJ2
JEJNJEJa
j=a
10|la| u m
+1
20n
3
(A
+1)
2n
205
v < |u11
It follows that
Ep W(
E
allu|12(1
|J11
a1|11111
[
+1
+ (A + 1) +
20n
3v 2np
40m
2n
20c
|||
a||x||2t 1
1
2
A p1 +1
20
6/P
(A.72)
(1
7)
20m
20
2
a|||u] a||x||2t
1
2
2ny ne log(n) +
20
a ||| 2
2
(1
20n 10c
n
V2nCn
80n 10c
for Cn when C is large enough. Here the second inequality uses |J 2 mn 10s; the third inequality uses (B3) in Lemma 4.1 and Assumption (A1); and the last inequality uses & < 0.01. By (A.69), it follows that [[W(')] || S 30t 30t vp/n. Thus we have
E E laf W
||(0|]
240./mn200
This lower bound for the normalized margin can be easily extended to the other v's. Applying Lemma A.9, have
cn 1-0 ||111*)
exp ( (-1() 52( n 20c ||lll ))
Plaay)~Paa (y # sgn(gg(( (s;; (((;(()()))) < 2exp
2402p
p
Lemma 4.6. Suppose that Assumptions (A1)-(A6) hold. Under a good run, we have that for
1 S Vn,
/ellull3 m
w (t)
+
Ω
t);
JEJ
1
Kun
(1) |121l =
aalla|1
vm +
a al|aa| MolllVlogie Mlg(n)) Vmn
1)
JEJa
Proof. This lemma essentially implied by the proof of Lemma 4.7. By (A.66), we know that for all j E Jn,
(w) +11) 0.
Then note that (w) +p1)) = ((w +u1)). From this we have
allaa12(1 - 22 - ((aillul1)) 40ym
(1))) +111
Σ
s((1 +wr))
[Tr
vmm
JUEJ
J:EJ
35

---

Published as a conference paper at ICLR 2024
where the first inequality comes from (A.67). Recall that in Lemma A.7, Aps is defined as |1++2 n n-wil pill + vn. Applying (B3) in Lemma 4.1, we have
n+m - - 7 + +c++)]++ (n( + C+m - -n/4)l
+ (m(n- CC-m1 nn/4] + |n- 7[n-- c-w1)l
<4v en enlloog)))
Then Ar is upper bounded by
Ap. S vn 4Venlog(n) (
Combining the inequality above with equation (A.70), we have
al|aal2 2n/m m
( vm
/mn allaal allull&V/oo(n)) ylog(n)))
Kuy -1)1
2)] 2]] in 1-+ (App +] [0(
+
vmn
JEJJ
A.5.3 PROOF OF THEOREM A.13: 1-STEP TEST ACCURACY
Before stating the proof, we begin with the necessary definitions and a preliminary result. Recall that (1) 9, (t) 1/2 and the decomposition (A.28). When = 0, we denote
(1) 5. =
(0)
aO
(0) aa) usa j E [m]
(A.73)
2n
and W( = [w]
Next lemma shows that W(() is a good approximation (())
w m,
with large probability.
Lemma A.10. Suppose Assumptions (AI) and (A2) hold. Given {x.} E Gdato and W(0) E Gww we have
(0) I phinit v 3m//2
OWwimip33 3m
W(
Σ
11 my 112
(1)
Vn
Proof. Let 2 (xr;; W(O)). Note that ('(x) = 11((1 exp(2)), we have [((==) 1/2 s |2|/2. 2. It yields that
|1
S
1 | 10
100
S
Σ 1=1
lay (we (0) xxd) S
m
w
(0)
1|
||0||
1=1
j=1
(A.74)
1 ] 1|
||ara|| |l S 2 purnin v v3m,
where the first inequality uses h(th) g. 1/2 and g (t) := (<) the second inequality uses triangle inequality; the third inequality uses Mauchy-Schwar. inequality: and the last inequality uses
(B1) in Lemma 4.1 and (C1) in Lemma 4.3. Denote hmax maxxe[m] [1: (()). Then have
(1) w 5 w 00
(0) o'( ((( h
(0) 244)))1a
n m
alimax
|| 2| + n(n - 1) max [x,] xn
n/m
admaa
v /3cummppp2 2/2
< n
/App <
Vni
36

---

Published as a conference paper at ICLR 2024
where the second inequality uses |1, 2p and p 2 Cn2Ilu|3 which come from (B1) and (B2) in Lemma 4.1 and Assumption (A2) respectively, and the third inequality uses (A.74). Further we have
Otwinie /3m
|| - 1 P
Σ
|0
wj 13
Vn
Lemma A.11. Suppose that Assumptions (AA))(A)) hold. Given X Goas for each [m]] have
n/24 S Var(D (0)
tmmg j) S n/2:
E (0) + - E (0)
+ . )1 )1 - n//22
Proof.
Recall that Ai C+1 UN- u1 A2 = C-rs UN+ According to equation (A.18), we have
D ( +unnd
I( 0). Σ 0) Σ
(A.75)
iEA2
According to Lemma A.15, we have
Var(D ) EBlfi (b,)).,,,,,,] 2 _ E [fi ... 1b,]]
1 Var Σ
Σ SEA2
b1)
|Ail + |A21
n 24'
2
iEA:
where fi ... := b, - -EEA,, by - (22l)/2)
2 0, and 69 are i.d Beroulli random variables defined in Lemma A.15, and the last inequality from (A.17). On the other side, similarly we have
Var D 1 S 2E IE ... [.,] = ((Ai) + 2
(A.76)
where the last inequality is from (B3) in Lemma 4.1. Denote f2 bb)) = (EiEA, b,
EiEr, b, (|Ail - (A2))/24 4 2 0, then have
E D -.
E[D (0)
| Eb[fe(b1,... bn)] S 2Eg [fa(b], ... bb,11
+p1
2E [[ Σ
(1/ 1/2
1
Σ (1/ }
ieAi
iEAr
(A.77)
< 16EE [[
Σ (b/ (W--111 }
+ Σ (1/ -
sEdo
< <4(|A,12 |A P) < n²
where the first inequality Lemma A.15; the second inequality uses (a + b)* < 8(a* b*); the third inequality uses the formula of the fourth central moment a binomial distribution with parameter equal to 1/2, i.e. (B 1/2)) n21 + (3n 6)/4)/4 < n? /4; and the last inequality from (B3) in Lemma 4.1. Combining (A.76) and (A.77), we have
D ( EDD 01 //S/V/V V Var(D - " J D
an mn E(D(O) (0) +m J 3/2
by applying the Cauchy-Schward inequality.
Lemma A.12. Suppose that Assumptions (A1)-(A6) hold. Given X = [Z1, xnn] T E Gant have
P(
Σ
m
(a,
3)
(0)
1
>
< 20
t tym C
3C,, Vne
Vm
37

---

Published as a conference paper at ICLR 2024
(ym C
P(
Σ ag a D 00)
+aivl
> >1) + 220(
30n Vnn vm
Proof. In this proof, by convention all PP),, B[.], Var()),p()) are implicitly conditioned on fixed X. Denote the expectation of D40 + by 6+p11 Note that conditioning X, (e are
i..d, and the expectation of ( +pr.. is
C+u = =C+tp - R4p - C- + nn-p)//2 2C, Vne,
(A.78)
where the inequality uses (B3) Lemma 4.1. By Lemma A.1I, we have
Var (D - z' p(D (m S 3/322
(A.79)
Denote
0 Var(ma,0la D f)) Ptu = pla
Combining (A.79) and results in Lemma A.14, we have
Eimayo(a,DD. ) 2
C+m, max n
tm
S o
max
" ft }: p+m S 32 max letmm 13..
48"
4
2' 2
(A.80)
Applying Berry-Esseen theorem, have
P( Σ a La
(ins)) c+m1 > f) S
2Φ
t/m
m
CnEPtM
S 2Φ Vm Vni + 20, /ne
tym
+ C'
O+p1
Vm
for some universal constant > 0. Here the second inequality uses o I (VTi+ + |e+ 1)? which comes from (A.80), and the last inequality uses (A.78). By the symmetry of aj. have
Emmaala, D f = O; Var(ma, la (0) D
(D ( p(ma,la, = E (0)
t 1
By (A.79), we have
24 + +M
$ E(D
n tu ) 2
te
E D +ui
$1 11
-un ) 13) S 8(n3/2 (A.81)
le+us 13).
Similarly, applying Berre-Esseen theorem, we have
m P
l/m
C
Σ
a Dyla,D mn
> 1) vn + 2C, / nE
2
mm'
where the inequality uses Var(ma, |a [0) +pn.j 1)) I (vn + letps 2) and (A.78). Then the results of this lemma proved by noting that Cn 2 1 for large enough
Theorem A.13. Suppose that Assumptions (AI)-(66 hold. With probability at least 1 1 ((1 mm) O(n=²) over the initialization of the weights and the generation of training data, after one iteration, the classifier sgn(f(r, W()) exhibits a generalization risk with the following bounds:
(1 - < IP (z.g)>Pam (y + szn(f(x; (())) $ (1(++
Proof. For any given training data X E Gamas denote the expectation of (00) by ev, i.e.
:= <=:=E(D X] (cv ny c-v + m22 (+m, +p2)
(A.82)
38

---

Published as a conference paper at ICLR 2024
and set of parameters 9x:
9x == {(a, W()) :
Σ a aoo (2)) - cw/21 < 3C, ne/m////mm//m
Σ
as ,ylay Llla D2 - 3C
GA, W( E gw]
Applying the union bound, we have
2C
x E Gdama 2 111-exp(--0 11-eep(_((lgg(mm
n
m
by Lemma A.12 and 4.3, Further we have
P((a, (()) E 9x,X E Gdataa 2 P(9x|X E Idma) E Gdata)
2C
2 - exp(- 1op(-Log''p (m))/2
-2n
vm
3C
2n-
Vm
Define events Fiest,y for test data:
Frest,y = fo E RP =||1=112 -p p- |||||31 &
(r, (V, S (01/plog(r) for all E [n]]. vE +/2)
Treat {ar U {=r in Lemma 4.1, we can show that P PaNN(e.1,)( E Fuall E Gdmaa 2 n= where Frest == Uvefda [up) } Fist,.. And Fust is symmetric set for x, i.c., if x E Fress then - also belongs to Fiest In the remaining proof, by convention all probabilities and expectations are implicitly conditioned on fixed X E Gaat and a, ( E Gx. Therefore, to simplify notation, we write IP(() and [[[] to denote IP( lla W()) (mx)) and E. la, (')), {ar; respectively. In other words, the fixed initialization and training data. We
new 'training' set with
datapoints. Following the proof procedure
randomness is over the test data (Ir, y), conditioned first look at the clusters centered at =u1 i.c. ~ N (( 4p), y = 1. Then we have
P P by + sgn(f(#, (() P N (f(z)) w()) 5)) M (()) < 0) 1 1 w()) > z).
(A.83)
2
Note that given (0) and X, we have with probability 1 that
m
I0,; S(x; w
Σ
a (ww
(2) ,,,))))))))))))))))) - 6b(11, (1) w (0) +للد|
j=1
(A.84)
<
la (meg
x>حl x
| (0) 12. |0:12
||
||ا||| S Wini v 3mmp /|||=]|
where the first inequality comes from the 1-Lipschitz continuity of ((-); the second inequality uses Cauchy-Schwarz inequality; and the last inequality uses Lemma 4.3. Next, recall that Wr is defined
39

---

Published as a conference paper at ICLR 2024
as in (A.73). By the same argument above, we have
(s;;; W - :()) - [(x; w W(
m
Σ
ag (0(( 9) w/z (0) ,))) o((w)) (1)
wj x\)]
(0)
j=1
las (w) W JT, x>|l 1
Σ
W 3
(1)
|0_| || - W
|1 ||=|||
j=1
j=1
I Sowiiiy 3mp/ |||a>|| I Wiit /33pp//laa
(A.85)
where the first inequality comes from the 1-Lipschitz continuity of Φ()); ); the second inequality uses Cauchy-Schwan inequality; the third inequality uses Lemma A0 and the last inequality uses Assumption (A3). Using (A.84) and (A.85), have by the triangle inequality that
(x; w() - f2; w w(0))] s 2wimit /mpla Ex
that for any IRP. (A.86)
Recall that
(w J w) x)) adj n
(1)
2n Σ
6'((w))
Then under good run, for Frest, have that with probability
(w wj x)
(1) (0)
adj
((0)
|ull
-CoVD
2m
Vm
where Cn = log(l))()) and the inequality uses the definition of Frest- It yields that
[(;; - w(o) - Σ
aaz
2n o(a,
(
th |u4l s aCnvt.
(A.87)
1=1
According to the definition of Gx. have
al|ual?
auj
ó(a, (
Mall
1 e-m1
||a|12
(A.88)
2n
4n
2//mn
Combining (A.86)-(A.88), we have
a||A|1
3a V (llg(m) 2/mm
((x:
4n e+m S + aCCVpt
(A.89)
The above inequality immediately implies that
[(f(æ; w(2) < o) 2 IP( a|lwal?
3aCn VE /elogm))
2n
e+u,
- aCC
2/mm
|
(A.90)
Similar to (A.89), for -I ~ N(- 1111T) Ip),
we have
s(=x;W(1)) al|1a|2
3aCn, Ve log(m)
2n c-m
S Sc
||
2/mn
Note that by definition, c-#1 -+tr the above inequality immediately implies that
a||/||2
3a 3aCn BaCnvElog(m)) 2/mm
P( w(1) FiFie)) 2 IP(
et ex aC v
2n
(A.91)
40

---

Published as a conference paper at ICLR 2024
According to the definition of Grest, have Ex I 4w v mmpp/ According to the definition of Gdata we have
Ic - The - C-u n 2 lcu - c-ul - In - -- 2 lcu + nle - C-u - n- 2/ Te-
2 (1 2n)m 2 n2122-// 12.
Thus we have le+ail
2 n 1/-5 2/. /4. yields that
0:
|
2
3a V log(m)
|+ppl - Ex aCn /
||||2
2n
2/mm
a||/1|1
- -4/mmp88/
Winit
Cn
up
30,, V log(m)
(A.92)
Vn
a||1|12
V
||||
2/m
a a||111?
Cn
3Cn
2VCno.01 > 00,
Vn
8n
mvn
3.n
where the first inequality uses |e- 2 n 2///22/44 and er S Auinity Auinity /mp3/2 the second inequality
uses Assumption (A5), (Al) and (A6); and the last inequality uses large enough. Combining (A.90)-(A.92), we have
(([(;;; W()) S F(ea)) + P(J(- W S OFFiea)
(P( 0111112 2n le+mll 6z aC vp+
3aCn V (CCVVllg(m)
(A.93)
[lel)) =1
2/mn
where the inequality uses Eg 2 0. Following a similar procedure, for the other side, we have
P(f(x; W()) - Ff P(f(-x; W()) - Fre
<P( a||1||? |e+pal 2 Ex aCnvp 3aC,, ve llo(mm)
(A.94)
|| (Fest))
2n
2/mn
Combining (A.93) and (A.94), have
((f W()) S Fleee) + P(f(-2; W()) S F(eu) =1
Following the same procedure, we have that for any +
(ENNGoL))Uff;; w()) < Fres) Po..N w()) S OfF =
Then for (z, y) ~ Potcan we have
1
P(a.a) eom w(D) P(y S F
2 7(1 n2 n=)
Ple (f(x:W(1) 5) < Pr (()((x))))))) s (ee P(F S 2 (1+
1
A.6 PROBABILITY LEMMAS
Lemma A.14. Suppose we have a random variable g that has finite La norm and a Rademacher variable a that is independent with g. Then we have
1 max! Var(g), [kg))) < Var(ap(ag))
(ar(()) (E(g))
(A.95)
[(aa(aag) - E((ab(ag))) 32 max [Elg Elg] ||[])) [[(]]
(A.96)
Proof. The expectation of the random variable ag(ag) is
Elad(ag)] = -((o(o) - o(-g)] E[g],
(A.97)
41

---

Published as a conference paper at ICLR 2024
where the first equation uses the law of expectation, and the second equation uses Φ(x) - -d(-x) The second moment of ao(ag) is
[((aa(ag))) )) H[(((g))) - + -1 -
(A.98)
where the last equation uses o(x)² + x(~x)² = x². Combining (A.97) and (A.98), we have Var(ac(ag)) -
2 |60 +
4 (واي)) Var(g) 4 (k[))))
which implies (A.95). Moreover, for a random variable X that has finite L3 norm, we have
| E[X]|| 13 < |||||3 + ||[ - |||||s + =|]j]]
|||
where the second inequality is due to |E X l3 ||x Thus we have
E X and the last inequality is due to || -
1 | | I (g]la 4 + -((-)) 41 "1
E ((((aa)))
2
where the last equation is due x(x)² x(-=)) |2/| Then by ||al|s I l1g E[gllls + Eiglt. we have
E ablag) [lol1" < 4 - Elgllls + [Bbl]) I 32 max - Elgl" El
Lemma A.15. Suppose 7. [21, ;n] N(0, Σ), where En; =1, and |Enzl S $
i * j S n. And Z' = 1,),,.. , 2/1 N(, In). Let b; I(2; > 0) and By [[ > 0), [m] be
Bernoulli random variables. Let B [b_.... bn] and B' = 1b. b/
Then we have that for
any nen-tegative function f R + R$ U (0).
EB-IS(K ... bb)1 s (bls(bi, ... bn)] S 2E [f(b], ... (f(6,,,..,,,,...
Proof. Note that for any fixed value (b), (b)) E [0,1)",PPp Pb ... ,,,)) = (1/2)". Then we have
EBlf(bl, bn)]
Σ
(b.....,,,))PPEE
bn)
by An
2 (2 y"
Σ
f(bi. ... ...,,,,, Pr (b1, ... .un)
(A.99)
= (2 [fhh.....
where the inequality comes from Lemma A.16. On the other side, similarly we have
Ebi b..)l < (B ,b)). By C 8, we have (2 )2 (1 (1-4/((CCC 4/(n))) 4/(CCn) 1/2 and (272)* (1 (1+4/(Cn))) s
(A.100)
exp(4/C) & exp(1 (2 $ 2. Combining these results with (A.99) and (A.100), we have
EB [f(b, bb)l s EBlf(b.... bn)] s 2E [f(b,. .,,,,,,,,))..
Lemma A.16. Suppose 7= (2 =n] T N(0, E), where En = 1, and |En1 S 1/(Cn)2) 15
i { j E n. Then we have that for any subset A E [n].
I(2 > 0)
II
1(4 0)] S 5
ieA
{E[R]AA
form 1/2 2/(Cn) and 72 1/2 + 2/(Cn).
42

---

Published as a conference paper at ICLR 2024
Proof. We first prove the result for A = [m]. Note that
P(z1 > 0, :.. ,n > 0) = P( 0) II
> P(=>0/= > 0, "0 > 0).
(A.101)
2=2
Let Zz-1 = --,, , 1 and denote the covariance matrix of 1.,,-.., ; as 2h-11
1 Ek
C
1
where - = Cov(Zx-1) and EN = CovCCoo/iaa (Z =k). Then |exs| S 1/((22) for E [k 1], and
the conditional distribution of 2 - is (e, Σ Z - c -x (k). theorem, have
(6). By Gershgorin circle
1
Amin EEx- x Amax (ER-1) 11
On
Cn
Denote fx-1 (-) as the density function of Z 1- Then we have
/0 ... .. fk-1(Zx-1)0
{ k-1 Zk-1
Σ-1
P(z+ > 0/2x-1 > .....,,.. >1>00 .2100
d2g d2k-1
0
Ek 5 1
--
ck kー k- Z
2
fw- (Z
d31
dzk
|0
1/3
1/2
Th 1522/0
C D k- Ck
1
| Ex -2
P( 1/2 1,200-1-1111111 2x-1 S 22
2
V 22m(1 - c) (x
2V2
2 2
)P(11 [111- 11-22-111 22/
11/2
nCva
1 nC
1 2v2
(1
3 On
1 2
(A.102)
for sufficiently large n. Here the second inequality uses (0(x) (()) < ('(0)az and Cauchy- Schwarz inequality; the third inequality uses mmin ( = Amin (Ex-1) 2 1/2 and | cxll S /2[|ee/| I //2m3/2/CC and the fourth inequality uses the concentration inequality for chi-square random variables in Lemma A.17. Then the result is proved by combining (A. and (0.).2). On the other side, we have
c/ 2
Zp-1
Pizx > 012k-1 > O,... 21 > 0) <
f 14-1(2x-1)4
dēg
dzk
11
1/2 152v/
-- ck
+ 1P12 k-1
1/2
14-1ll > 22/7)
(1
1
| <x 22V7
+ +P(11E///222 Zk-1ll > 2/n)
2
2 c &T2iLLEN)
1
2
1 2
+ cxp(- S
2
nC
On'
Note that our proof does not use any information related to A, thus we can extend the result for any subset C [r].
Lemma A.17. For X2 ii.d ~ N(0, S we have
(0)/L>P((())) to) < exp( - 2212), 14 2 1;
43

---

Published as a conference paper at ICLR 2024
7 x2 1 2 <
P(
2exp(-n1218))
14 E (0. 1).
na²
k=1
Proof. For the first inequality, we note that
+000
exp( x² 2²²)dx² (1)
((t)=
VTT x
exp(
x²)dr dd 1
2
yields that for any t 1, On the other side, we have
P([X) to) 2Tt(t) S 2t/ S exp(-t3/22)
12
exp( (
+00
(t) 2
2n 12 exp(-
22 V2T 1 +
(())).
x
z=t
1 1 t
When 2 further yields that (1) 2 (2)). Thus we have
P([X1 2 to) (t) 2 4 (t)/t.
The second inequality is Example 2.11 in Wainwright (2019)
Lemma A.18 (Hoeffding's inequality, Equation (2.11) in Wainwright (2019)). Let X 1 s be a series of independent random variables with X E [a, 6]. Then
Xx (X+ 2 exp (- n(b a) V1
212
Lemma A.19. [Berry-Essed Theorem, Theorem 3.4.17 Durret (2019)) Let XI,... , Xx, are
i.dd random variables with [[[,,] = O, Var(X;) = 0², and x; p < co. If Fn(x)) is the distribution of E / /n), then
Xa/(o
3p
[Fx(2) ((12)-4(11)) x(x) S
σ nn
A.7 EXPERIMENTAL DETAILS
In experiments, dimension 40000, number of train/test samples 200 μ= 2.5 p/n, number of neurons 1000, label noise rate 0.05, and initial weight scale Winit = 10 15 For
Figure 3, 2, and I-left, the step size a = 10 12 For Figure 4 and I-right, a = 10 166
0
25
100
1000
2000
5000
Number of neurons (Frequency)
0.02
0.02
0.02
0.02
0.02
0.02
0.05
0.02
0.03
0.20
0.25
0.25
(w//0) jua)
-1
-2
-2.5
0.0
2.5
-5
(w/(),u))
1e-13
1e-13
1e-13
1e-12
1e-12
1e-12
Figure 4: Histograms of inner products between positive neurons and j pooled over 100 independent runs under the same setting as in Figure but with smaller step size. Top (resp. bottom) row: Inner products between positive neurons and μ1 (resp. H2). While the projections of positive neurons ((() onto the μa and μ2 directions have nearly the same distribution when the network cannot generalize, they become much more aligned with =yn when the network generalize.
44

---

Published as a conference paper at ICLR 2024
A.8 ADDITIONAL EXPERIMENTS
In this section, we show additional experiments with setups that are variations on that of Figure 1. See Appendix A.7 for details of the setup of the experiments in Figure 1.
1.0
0.9
accuracy
0.8
train (med)
0.7
test (med) 20-80 qtl
0.6
acc=0.5
0.5
100
10a
102
steps
Figure 5: More samples hurt generalization. Same experimental setup in Figure (Left) except that the sample size n 2000 is 10 times larger. Even running over 10x number of steps, the median test accuracy remains near 90%, while in Figure 1 the median test accuracy is 100%.
1.0
accuracy
0.8
init_scale 1e-14
init Scale 1e-13
109
10a
10a
1.0
1.0
accurecy
0.8
Int scale 18-12
0.8
scale
1e-11
0.6
10a
10
107
10a
steps
steps
Figure 6: Larger initial weight scale. Same experimental setup in Figure 1 except that the initial weight scale Wint are larger. (Recall that Wini 10-15 in Figure 1). Moreover, the experiments corresponding winit 10 (resp. Winit = 10 10-11) are ran with 10x (resp. 100x) number of steps.
sccuracy
0.8
0.8
(med) test (med)
a.c-0.5
Figure 7: Train both layers. Same experimental setup as in Figure 1 except that the outer layer weights are also trained. The accuracy dynamics is similar that of Figure 1.
45


Reviews:
[{"review_id": "9itkvJ0ZkG", "reviewer": "Reviewer_6CDP", "summary": "This paper analyzes two-layer ReLU network trained by gradient descent on XOR cluster data with a high-dimensional input space, and rigorously proves grokking and benign overfitting occurs.\n\nSpecifically, with one sufficiently large gradient step, the network is almost a linear classifier and achieves perfect overfitting to the training data, which contains label-flipping noise. If training is continued, the model almost perfectly predicts the clean label, while keeping perfect overfit to the training data.", "questions": "- Is it possible to see an additional experiment when the initialization scale is not so small?\n\n- When $\\|\\mu\\|$ gets small, is this grokking phenomena still observed?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "### Notable results on benign overfitting of neural networks beyond linearly separable data\n\nAs far as I understand, proving benign overfitting for neural network involves several difficulties due to its nonlinearity, and especially I agree that showing the superiority of neural network to linear methods by learning nonlinear target function has been largely open in this context. I think XOR cluster data is a good starting point to this problem and this paper proves benign overfitting under moderate assumptions.\n\n### Providing useful theoretical understandings on grokking\n\nAs well as benign overfitting result, this paper also proves that after one large (compared to the initialization scale) gradient step, the neural network approximately behaves as the linear model and perfectly overfit to the training data. The phenomena that the early stage of training only produces the linear model but there is a transition to the nonlinear neural network with richer features is particularly interesting, providing one of the first rigorous theoretical demonstrations of grokking.\n\n### The paper is well written and clearly explaining its theoretical key points.\n\nThe proof sketch section is well-written and provides a sufficient understanding of the overall theoretical contributions. It is expected that the techniques presented in this paper will also be valuable in demonstrating the potential for more enriching feature learning in the future.", "weaknesses": "### Justification of the small initialization\n\nIn my understanding, it is crucial to take a small initialization scale compared to the step size $\\alpha$ to obtain the perfect overfitting at the first gradient step. I think this is acceptable as theory, but it should be better to justify such a small initialization. I also want to know what happens if the initialization scale is much larger than used in Figure 3.\n\n### Large signal-to-noise ratio\n\nCompared to [Ji & Telgarsky (2019)](https://arxiv.org/abs/1909.12292); [Wei et al. (2019)](https://arxiv.org/abs/1810.05369); [Barak et al. (2022)](https://arxiv.org/abs/2207.08799); [Telgarsky (2023)](https://openreview.net/forum?id=swEskiem99); [Suzuki et al. (2023)](https://openreview.net/forum?id=tj86aGVNb3&referrer=%5BAuthor%20Console%5D(%2Fgroup%3Fid%3DNeurIPS.cc%2F2023%2FConference%2FAuthors%23your-submissions), where $x$ is a $d$-dimensional (essentially) rotationally invariant input and $y=\\rm{sgn}(x_1x_2)$, this paper considers large signal $\\|\\mu\\|$ as an input.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2KLOTeJSAm", "reviewer": "Reviewer_mTpn", "summary": "This paper delves into the exploration of benign overfitting and \"grokking\" in two-layer ReLU neural networks when trained on XOR cluster data with noisy labels. The authors demonstrate that networks can perfectly fit noisy training data (benign overfitting) and, after a period, transition from harmful overfitting to a stage where they generalize near-optimally (\"grokking\"). Through rigorous theoretical analysis and proofs, the study reveals that these surprising phenomena are evident in the networks' training trajectories, providing a nuanced understanding of overfitting and generalization in neural network models.\n\nThe paper's contributions lie in providing the first theoretical insights into benign overfitting in non-linearly separable data distributions, unraveling the feature learning dynamics under gradient descent. These findings illuminate the pathways through which neural networks navigate the complexities of noisy data, offering a fresh perspective on their capacity to generalize despite apparent overfitting.", "questions": "* Could you provide more insights or justification regarding Assumption A1? Specifically, could you clarify why an increase in the number of training samples seems to negatively influence the model, contrary to the common understanding that more samples generally improve a model's generalization capability?\n\n* Could you elaborate on the mechanism of overfitting after the neural network learns the directions of $u_1$ and $u_2$.\n\n* Could you clarify the role and significance of Lemma 4.6 in the context of the paper's objectives and findings?\n\n* Could you make a comparsion your results and techniques with [1] and [2]?\n\n* Have you considered variations in the network architecture, such as not fixing the second layer, and if so, how do these variations influence the results?\n\n* In the model settings used in the paper, there doesn’t appear to be an upper bound on the network width. If an extremely wide network setting were used, would the findings align with the \"lazy training\" regime? Could you discuss how the results might be influenced by varying the width of the network to such extremes?\n\nI would increase my score if the authors could clarify my concerns demonstrated in the above questions.\n\n--------------------------------------\nI increase my score to 6 after rebuttal.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper offers a new theoretical examination of benign overfitting and \"grokking\" in two-layer ReLU neural networks. It focuses on XOR cluster data with noisy labels, giving a detailed exploration and proofs related to these phenomena. The authors use existing concepts and new theories to better explain the behavior of neural networks with noisy training data.\n\n* The paper is structured and clear, effectively communicating the authors’ work and results. It has a logical organization that makes it easy for readers to follow the ideas and analyses. The presentation of definitions, explanations, and proofs is mostly straightforward, helping readers understand the complex concepts and findings.\n\n* The paper is important because it helps understand overfitting and generalization in neural networks better. It explains the concepts of benign overfitting and \"grokking\" in one framework.", "weaknesses": "* The assumption made in A1 seems to contradict common understanding. Generally, increasing the number of samples, even with limited noisy labels, tends to enhance the generalization capability of neural networks. However, in Assumption A1, having a larger number of training samples seems to adversely affect the model, as indicated by its presence on the right-hand side of the inequality. This aspect might require further clarification or justification within the context of the study.\n\n* The mechanism of overfitting, once the neural network learns the directions of $u_1$ and $u_2$. is not explicitly clear. The paper mentions that post the initial gradient step, positive neurons learn $u_1$ while negative neurons learn $u_2$. However, the explanation seems lacking in how the network overfits to samples with noisy (flipped) labels in this condition. A more detailed discussion or clarification on this would be beneficial.\n\n* The role of Lemma 4.6 in the paper is unclear. A clearer explanation of how it relates to other parts of the paper and how it contributes to the overall arguments and conclusions is needed for better understanding. \n\n* The paper seems to lack a comparative discussion with some relevant works, specifically references [1,2].\n\n\n[1] Meng, Xuran, Difan Zou, and Yuan Cao. \"Benign Overfitting in Two-Layer ReLU Convolutional Neural Networks for XOR Data.\" arXiv preprint arXiv:2310.01975 (2023).\n\n[2] Glasgow, Margalit. \"SGD Finds then Tunes Features in Two-Layer Neural Networks with near-Optimal Sample Complexity: A Case Study in the XOR problem.\" arXiv preprint arXiv:2309.15111 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Q5gdGK9bau", "reviewer": "Reviewer_AQUA", "summary": "This paper theoretically combines the phenomena of benign overfitting and Grokking by training a two-layer ReLU neural network using gradient descent on XOR cluster data. The authors demonstrate that after the first step of gradient descent, the network achieves 100% training accuracy, perfectly fitting the noisy labels in the training data, but exhibits nearly random performance on the test data. However, after some time in training, the Grokking phenomenon occurs, and the network achieves near-optimal test accuracy while still adapting to the random labels in the training data, demonstrating benign overfitting.", "questions": "1: Do the authors have any ideas for improvements in the high-dimensional setting?\n\n2: What behavior does the test error exhibit when the training time is between 1 and $n^{0.01}$?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper is the first to study the combination of benign overfitting and the Grokking phenomenon in neural networks.", "weaknesses": "My main concern about this paper lies in its assumptions. Combining assumptions A1 and A2, we can obtain $p\\geq C^4 n^{5.02}$. This is an extremely high-dimensional setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)