SGD Finds then Tunes Features in Two-Layer Neural Networks with near-Optimal Sample Complexity: A Case Study in the XOR problem

TL;DR

This paper demonstrates that stochastic gradient descent can efficiently learn the XOR function in a two-layer neural network with near-optimal sample complexity, revealing a two-phase feature learning process.

Contribution

It provides the first analysis showing SGD's ability to learn XOR with polylogarithmic samples, highlighting a two-phase feature evolution in training.

Findings

SGD achieves near-optimal sample complexity for XOR learning.

Network evolves through signal-finding and signal-heavy phases.

Training only a small fraction of neurons suffices for feature amplification.

Abstract

In this work, we consider the optimization process of minibatch stochastic gradient descent (SGD) on a 2-layer neural network with data separated by a quadratic ground truth function. We prove that with data drawn from the $d$-dimensional Boolean hypercube labeled by the quadratic ``XOR'' function $y = -x_ix_j$, it is possible to train to a population error $o(1)$ with $d \:\text{polylog}(d)$ samples. Our result considers simultaneously training both layers of the two-layer-neural network with ReLU activations via standard minibatch SGD on the logistic loss. To our knowledge, this work is the first to give a sample complexity of $\tilde{O}(d)$ for efficiently learning the XOR function on isotropic data on a standard neural network with standard training. Our main technique is showing that the network evolves in two phases: a $\textit{signal-finding}$ phase where the network is small and many of the neurons evolve independently to find features, and a $\textit{signal-heavy}$ phase, where SGD maintains and balances the features. We leverage the simultaneous training of the layers to show that it is sufficient for only a small fraction of the neurons to learn features, since those neurons will be amplified by the simultaneous growth of their second layer weights.