Learning threshold neurons via the "edge of stability"

TL;DR

This paper investigates the 'edge of stability' phenomenon in neural network training with large learning rates, revealing a phase transition affecting the learning of threshold neurons and its implications for generalization.

Contribution

It provides a theoretical analysis of the edge of stability in simplified neural models, showing a phase transition that influences threshold neuron learning and generalization.

Findings

Established the edge of stability phenomenon in simplified models.

Discovered a phase transition for step size affecting threshold neuron learning.

Suggested a link between the edge of stability and improved generalization.

Abstract

Existing analyses of neural network training often operate under the unrealistic assumption of an extremely small learning rate. This lies in stark contrast to practical wisdom and empirical studies, such as the work of J. Cohen et al. (ICLR 2021), which exhibit startling new phenomena (the "edge of stability" or "unstable convergence") and potential benefits for generalization in the large learning rate regime. Despite a flurry of recent works on this topic, however, the latter effect is still poorly understood. In this paper, we take a step towards understanding genuinely non-convex training dynamics with large learning rates by performing a detailed analysis of gradient descent for simplified models of two-layer neural networks. For these models, we provably establish the edge of stability phenomenon and discover a sharp phase transition for the step size below which the neural network fails to learn "threshold-like" neurons (i.e., neurons with a non-zero first-layer bias). This elucidates one possible mechanism by which the edge of stability can in fact lead to better generalization, as threshold neurons are basic building blocks with useful inductive bias for many tasks.