Source-linked AI summary
Benign Overfitting and Grokking in ReLU Networks for XOR Cluster Data
Zhiwei Xu, Yutong Wang, Spencer Frei, Gal Vardi, Wei Hu
TL;DR
The paper asks whether benign overfitting and grokking can be proved for neural-network classification on non-linearly separable data, where existing benign-overfitting results were restricted to linearly separable distributions. It analyzes gradient descent on noisy XOR clusters and proves an initial catastrophic-fit phase followed by near-perfect test generalization while noisy training labels remain fit.
Problem
Theoretical benign-overfitting results for neural-network classification were restricted to linearly separable data, and prior work had not rigorously proved grokking in a neural-network setting.
Source: §1 · p. 2Method
The authors analyze feature learning along gradient descent in a two-layer ReLU network trained on XOR-cluster data with flipped labels.
Source: §1 · p. 2, §4.2.2 · p. 10Results
100% training accuracy at t=1 accompanies near-random test performance; for Cn^0.01 ≤ t ≤ √n, test error is exp(-Omega(n^2.01)) while training remains perfect.
Source: §3 · p. 6, §5 · p. 11Takeaways & Limitations
The result provides a theoretical characterization of benign overfitting for neural-network classification on a non-linearly separable distribution.
Source: §1 · p. 2The analysis bounds the number of training steps for technical reasons, leaving generalization at longer training times open.
Source: §3 · p. 7
Abstract
from arXiv · showhide
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.
1 Introduction
The paper establishes benign overfitting and grokking in a two-layer ReLU network on noisy XOR-cluster data, a non-linearly separable setting. It traces the transition from an initially non-generalizing fit to later generalization and situates this behavior against prior work.
1 Introduction
- Existing neural-network results on benign overfitting were restricted to linearly separable distributions, and prior work had not rigorously proved grokking in a neural-network setting.The paper frames its contribution around these two open gaps.Source: §1 · p. 2
- After one gradient-descent step, the network fits every training label but performs near chance on test data; later it reaches near-perfect test accuracy while retaining a perfect training fit.The authors identify this as a transition from catastrophic to benign overfitting, or grokking.Source: §1 · p. 2
- After one step the network behaves approximately like a linear classifier that fits the training data but cannot generalize; later its neurons align with cluster means ±µ1 and ±µ2.The feature alignment supplies the generalizable features needed for the XOR task.Source: §1 · p. 2
1.1 Additonal Related Work
- Prior neural-network benign-overfitting studies covered cluster and image-patch settings, but the cited trajectory-based analyses did not identify grokking.The related-work discussion also describes stationary-point analyses and a taxonomy of catastrophic, benign, and tempered overfitting.Source: §1.1 · p. 3
- Earlier grokking studies proposed mechanisms and examples, but the introduction says none had established a rigorous proof in a neural-network setting.The cited work spans transformers, linearly separable classification, and other grokking explanations.Source: §1 · p. 2, §1.1 · p. 3
- On the same noisy XOR-cluster setting, prior work used early stopping to obtain clean test accuracy with training error near the label-noise rate; this paper instead finds catastrophic overfitting before later benign overfitting.The authors describe the earlier early-stopping behavior as fundamentally different from their result.Source: §1 · p. 2, §1.1 · p. 4
2 Preliminaries
The preliminaries specify notation, the noisy XOR-cluster data distribution, and the two-layer ReLU network and gradient-descent training setup.
2.1 Notation
- The notation defines distribution and probability symbols, vector and matrix norms, cosine similarity, and asymptotic order conventions.It also defines {±µ} as the pair of vectors +µ and −µ.Source: §2.1 · p. 4
2.2 Data Generation Setting
- The data are sampled from four Gaussian clusters centered at ±µ1 and ±µ2, with labels flipped at rate η; the cluster means are assumed orthogonal for simplicity.The sample size is assumed sufficiently large, and the two mean vectors have equal norm.Source: §2.2 · pp. 4–5
2.3 Neural Network, Loss Function, and Training Procedure
- The model is a two-layer ReLU network trained by gradient descent on empirical logistic loss, with second-layer weights fixed during training.The first-layer weights are initialized randomly and updated using step size α.Source: §2.3 · p. 5
3 Main Results
Under assumptions on signal strength, dimension, noise, step size, initialization, and network width, the theorem shows that the network first fits noisy labels with near-random test performance, then generalizes while continuing to fit them. This gives a guarantee of benign overfitting for neural-network classification on a nonlinear data distribution.
3 Main Results
- The theorem assumes high signal-to-noise ratio and dimension, constant-bounded label noise, a small step size, suitable initialization scale, and at least Cn^0.02 neurons.The assumptions support near-orthogonal training examples and the concentration arguments used at initialization.Source: §3 · pp. 5–6
- 100% training accuracy at t = 1 pairs with near-random test error; for Cn^0.01 ≤ t ≤ √n, test error is exp(-Ω(n^2.01)) while training remains fit.The theorem holds with probability at least 1 − n^−Ω(1) − O(1/√m) over data generation and weight initialization.Source: §3 · p. 6
- Theorem 3.1 gives the first guarantee of benign overfitting in neural-network classification for a nonlinear data distribution.The stated result applies to the XOR cluster setting with a constant fraction of flipped training labels.Source: §3 · p. 6
- The guarantee is limited to t ≤ √n: controlling the loss ratio keeps flipped-label examples from having outsized influence, and longer XOR training is left open.The XOR dynamics require tighter ratio control than related settings discussed in the paper.Source: §3 · pp. 6–7
4 Proof Sketch
The proof sketch establishes high-probability properties of the data and initialization, then uses neuron dynamics to explain one-step overfitting and later generalization. In particular, neurons gradually align with cluster means, supporting generalization after the initial non-generalizing fit.
4.1 Properties of the Training Data and Random Initialization
- Under Assumptions (A1) and (A2), the sampled training data satisfy properties (B1)–(B4) with probability at least 1 − O(n^-ε).The high-dimensionality assumption yields near-orthogonal training points, a property used throughout the optimization and generalization proofs.Source: §4.1 · pp. 7–8
- With probability at least 1 − O(n^-ε), the initialization and data satisfy conditions that spread neuron activations and provide neurons aligned with cluster directions.Conditions (D1)–(D4) include broad activation across neurons and alignment with each cluster direction.Source: §4.1 · pp. 8–9
4.2.1 One-Step Catastrophic Overfitting
- At t = 1, the network behaves like a linear classifier that fits the training data but has 50% test error on the XOR distribution.Near-orthogonal training data allow the linear classifier to fit the samples, while opposing cluster means limit its test performance.Source: §4.2 · p. 9, §4.2.1 · p. 10
4.2.2 Multi-Step Generalization
- During later training, positive neurons align with ±µ1 and negative neurons with ±µ2, while each group forgets the opposing class directions.The proof sketch identifies this gradual feature learning as the mechanism that makes the network generalizable.Source: §4.2.2 · pp. 10–11
- For Cn^10ε ≤ t ≤ √n, Theorem 4.8 bounds generalization error as positive and negative neurons align with their respective class means.The proof links this alignment to neurons in the relevant groups dominating outputs for the four clusters.Source: §4.2.2 · pp. 10–11
5 Discussion
The discussion highlights a transition from early catastrophic overfitting to later generalization while the network continues fitting noisy labels. It also identifies limits on the proof’s time horizon and dimensionality assumptions.
5 Discussion
- The network continues to fit noisy training labels while later achieving near-zero test error, demonstrating grokking and benign overfitting together.The paper presents this as benign overfitting for neural-network classification on a non-linearly separable distribution.Source: §5 · p. 11
- Grokking occurs without explicit regularization in this setting, and the discussion attributes significance to gradient descent’s implicit regularization.The early classifier-like behavior fits high-dimensional training data but does not generalize on the XOR distribution; later, cluster-mean features support generalization.Source: §5 · p. 11
- The analysis bounds the number of training steps and assumes near-orthogonal data, requiring ambient dimension large relative to sample count.The generalization behavior as time tends to infinity and the effect of lower-dimensional data remain open questions.Source: §5 · p. 11
A.1 Additional Notation
This notation section defines standard normal distribution functions, common probability distributions, and two moments of a random variable.
A.1 Additional Notation
- The notation defines Φ and Φ′ as the standard normal CDF and PDF, ¯Φ as its upper tail, Bern(p) and B(n, p), and Var(X) and ρ(X) as variance and absolute third central moment.It also specifies the Bernoulli success probability and the binomial size and probability parameters.Source: §A.1 · p. 15
A.2.1 Proof of Lemma 4.1
Lemma 4.1 establishes that the sampled training data satisfy properties (B1)–(B4) with high probability. The proof bounds sample geometry and center counts using concentration arguments.
A.2.1 Proof of Lemma 4.1
- (B1) bounds each sample’s projections onto the centers and controls its squared norm; (B2) bounds pairwise inner-product deviations.The stated bounds are 10√log n∥µ∥ and 10√p log n for (B1), and 10√p log n for (B2).Source: §A.2.1 · p. 15
- (B3) controls center counts and noisy-label counts, while (B4) requires lower bounds on differences between opposite-center counts.The proof demonstrates representative count bounds and says the remaining inequalities follow analogously.Source: §A.2.1 · pp. 15, 18
- With probability at least 1 − O(n^-ε), the training data satisfy properties (B1)–(B4).The proof defines Gdata as the set of training data satisfying these properties.Source: §A.2.1 · pp. 15–16
- The proof obtains the geometric properties from Gaussian-noise bounds and union bounds, and handles center counts with Hoeffding and Berry–Esseen arguments.The data are modeled as Gaussian samples around one of the four centers.Source: §A.2.1 · pp. 16, 18
A.2.2 Proof of Corollary 4.2
Under assumptions (A1) and (A2), and conditions (B1) and (B2), Corollary 4.2 establishes near-orthogonality for distinct training samples.
A.2.2 Proof of Corollary 4.2
- Corollary 4.2 states near-orthogonality of the training data under assumptions (A1), (A2), and conditions (B1), (B2).Its conclusion applies to every pair of distinct sample indices.Source: §A.2.2 · p. 19
- The corollary’s stated scope is all pairs of sample indices i and k with i ≠ k.Source: §A.2.2 · p. 19
- The proof applies Lemma 4.1’s conditions (B1) and (B2) to distinct samples.The subsequent bounds use assumptions (A1) and (A2).Source: §A.2.2 · p. 19
A.3 Properties of the initial weights and activation patterns
This section introduces notation for neuron signs, activation patterns, and neurons aligned with the data centers, used in the proofs of Lemmas 4.3 and 4.4.
A.3 Properties of the initial weights and activation patterns
- The notation JP and JN partitions neurons according to whether their output weights aj are positive or negative.Source: §A.3 · p. 19
- A(t) denotes the sample–neuron pairs for which the neuron is active on the sample at time t.The text also introduces subsets Ai,(t) and A(t) as related activation-pattern notation.Source: §A.3 · p. 19
- The section defines sets of (ν, κ)-aligned neurons for each center ν and parameter κ, using n±ν := nν + n−ν.Source: §A.3 · pp. 19–20
A.3.1 Proof of Lemma 4.3
Lemma 4.3 gives high-probability properties of the random initialization under assumptions (A1), (A2), and (A6), including conditions on initial weights and output-weight signs.
A.3.1 Proof of Lemma 4.3
- With probability at least 1 − O(n^-ε), the random initialization satisfies the properties stated in Lemma 4.3.The lemma assumes (A1), (A2), and (A6).Source: §A.3.1 · p. 20
- The lemma groups initializations satisfying (C1) into GW and output-weight vectors satisfying (C2) into GA.The joint event W(0) ∈ GW and a ∈ GA has probability at least 1 − O(n^-ε).Source: §A.3.1 · p. 20
- For (C2), the proof uses that positive and negative output-weight counts each follow B(m, 1/2), applying Hoeffding’s inequality to bound their sizes.The final probability bound invokes assumption (A6).Source: §A.3.1 · p. 20
A.3.2 Proof of Lemma 4.4
This section develops probability and trajectory estimates for analyzing classifier behavior during training, including training-data fit and generalization bounds across specified steps. It also gives conditional-positivity bounds and reports experimental distinctions in positive-neuron projection alignment between failed and successful generalization.
A.3.2 Proof of Lemma 4.4
- With probability at least 1 − O(n^-ε), each sample activates a large proportion of positive and negative neurons, including aligned neurons of both signs.The proof also states that nearly all neurons in each sign group satisfy the corresponding property.Source: §A.3.2 · p. 20
A.3.3 Proof of the Probability bound of the “Good run” event
- The cited passage states that the authors combine the probability lower-bound parts of Lemmas 4.1, 4.3, and 4.4, without specifying the resulting bound.Source: §A.3.3 · p. 26
A.4 Trajectory Analysis of the Neurons
- The training dynamics are nearly unchanged from their first-step behavior for the first O(p) steps.The trajectory analysis states this under a good run, while Lemmas A.3 and A.4 provide controls over the interval 0 ≤ t ≤ 1/(√npα) − 2.Source: §A.4 · pp. 26–27
A.4.1 Proof of Lemma A.3
- Under a good run, Lemma A.3 bounds the maximum h_i(t) across samples throughout 0 ≤ t ≤ 1/(√npα) − 2.The supplied passage states the interval and maximum but does not expose the bound’s right-hand side.Source: §A.4.1 · p. 27
A.4.2 Proof of Lemma A.4
- Under a good run, Lemma A.4 gives a bound for each sample k over 0 ≤ t ≤ 1/(√npα) − 2.The lemma defines the clean label by cluster center: yν = 1 for ν ∈ {±µ1} and yν = −1 for ν ∈ {±µ2}.Source: §A.4.2 · p. 28
A.4.3 Proof of Corollary A.5
- Corollary A.5 states trajectory bounds for neuron–datapoint pairs according to the sign of a_j y_k, over time windows ending at 1/(√npα) − 2 or 1/(√npα) − 3.The supplied passages give the E1–E3 conditions but truncate parts of their inner-product inequalities.Source: §A.4.3 · pp. 30–31
A.4.4 Proof of Lemma A.6
- Lemma A.6 proves trajectory properties for neurons aligned with +µ1 over 1 ≤ t ≤ 1/(√npα) − 2.Its proof analyzes the cluster centered at +µ1 and treats neurons with positive and negative output weights separately.Source: §A.4.4 · pp. 31–32
A.4.5 Proof of Lemma A.7
- Lemma A.7 proves corresponding trajectory properties for neurons aligned with −µ1 over 2 ≤ t ≤ 1/(√npα) − 2.The proof uses the same data as Lemma A.6 while flipping the sign of a_j in the trajectory analysis.Source: §A.4.5 · p. 33
A.5.1 Proof of Theorem A.8: 1-step Overfitting
- The classifier correctly classifies all training datapoints for 1 ≤ t ≤ 1/(√npα) − 2.This result holds under Assumptions (A1)–(A6) and a good run.Source: §A.5.1 · pp. 37–38
A.5.2 Proof of Theorem 4.8: Generalization
- For Cn^10ε ≤ t ≤ √n, Theorem 4.8 establishes an upper bound on the classifier’s generalization error under a good run.The supplied passage does not include the bound’s expression.Source: §A.5.2 · p. 39
- The generalization proof extends a normalized-margin lower bound across ν and states Lemma 4.7 for good runs with 1≤t≤√n.The supplied passages do not include the bound’s explicit expression.Source: §A.5.2 · pp. 42–43
A.5.3 Proof of Theorem A.13: 1-step Test Accuracy
- After one iteration, clean test risk lies between 1/2(1−n^-ε) and 1/2(1+n^-ε), with probability at least 1−O(1/√m)−O(n^-ε).The probability is over weight initialization and training-data generation.Source: §A.5.3 · p. 46
A.6 Probability Lemmas
- For independent Rademacher a and finite-L3-norm g, Lemma A.14 bounds Var(aϕ(ag)) in terms of Var(g) and (E[g])².The displayed bound’s operators are garbled in the supplied passage, so its exact endpoints are omitted.Source: §A.6 · p. 50
- When Gaussian coordinates have unit variance and off-diagonal covariances at most 1/(Cn²), sign-indicator expectations for nonnegative f are within factors 1/2 and 2 of the independent case.The comparison is between correlated Gaussian sign indicators and independent Bernoulli(1/2) indicators.Source: §A.6 · p. 51
A.7 Experimental details
- When generalization succeeds, positive-neuron projections align more with ±µ1; when it fails, their distributions across the µ1 and µ2 directions are nearly the same.Figure 4 uses a smaller step size than Figure 1; the setup uses p = 40000, n = 200, m = 1000, η = 0.05, and ωinit = 10−15.Source: §A.7 · pp. 53–54, Figure 4 · p. 54