Benign Overfitting and Grokking in ReLU Networks for XOR Cluster Data

TL;DR

This paper provides the first theoretical proof of benign overfitting and grokking phenomena in two-layer ReLU neural networks trained with gradient descent on XOR cluster data with noisy labels, revealing the feature learning process.

Contribution

It demonstrates that both benign overfitting and grokking occur in neural networks on non-linearly separable data, with detailed analysis of feature learning over training.

Findings

Network fits noisy labels after first GD step with poor test performance.

Later training leads to near-optimal test accuracy while still fitting noise.

Feature learning transitions from linear to generalizable representations.

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.