Understanding Edge of Stability in Rank-1 Linear Models for Binary Classification

Recent research in deep learning optimization reveals that many neural network architectures trained using gradient descent with practical step sizes, $\eta$, exhibit an interesting phenomenon where the top eigenvalue of the Hessian of the loss function, ${\lambda }_1^H$, increases to and oscillates about the stability threshold, $\frac{2}{\eta }$. The two parts of the trajectory are referred to as Progressive Sharpening and Edge of Stability. The oscillation in ${\lambda }_1^H$ is accompanied by a non-monotonically decreasing training loss. In this work, we study the Edge of Stability phenomenon in a two-layer rank-$1$ linear model for the binary classification task with linearly separable data to minimize logistic loss. By capturing the core training dynamics of our model as a low-dimensional system, we rigorously prove that Edge of Stability behavior is not possible in the simplest one datapoint setting. With two datapoints, empirically we observe that, when the margin between the two points is small, Edge of Stability may occur. The oscillation in loss and sharpness becomes perpetual in the limit where the margin converges to 0. Using bifurcation and normal form analysis of the dynamical system in this asymptotic setting, we provide a sufficient condition under which the dynamical system exhibits perpetual oscillation in both loss and sharpness.