Learning Dynamics of Deep Linear Networks Beyond the Edge of Stability

TL;DR

This paper analyzes the learning dynamics of deep linear networks beyond the edge of stability, revealing loss oscillations, chaos, and the role of symmetry-breaking in training behavior.

Contribution

It provides a detailed theoretical analysis of deep linear networks' behavior beyond EOS, including loss oscillations, chaos, and the impact of symmetry-breaking on training dynamics.

Findings

Loss oscillations follow a period-doubling route to chaos.

Loss oscillations occur within a subspace characterized by the learning rate.

Symmetry-breaking at EOS leads to monotonic decay of the balancing gap.

Abstract

Deep neural networks trained using gradient descent with a fixed learning rate $\eta$ often operate in the regime of "edge of stability" (EOS), where the largest eigenvalue of the Hessian equilibrates about the stability threshold $2/\eta$. In this work, we present a fine-grained analysis of the learning dynamics of (deep) linear networks (DLNs) within the deep matrix factorization loss beyond EOS. For DLNs, loss oscillations beyond EOS follow a period-doubling route to chaos. We theoretically analyze the regime of the 2-period orbit and show that the loss oscillations occur within a small subspace, with the dimension of the subspace precisely characterized by the learning rate. The crux of our analysis lies in showing that the symmetry-induced conservation law for gradient flow, defined as the balancing gap among the singular values across layers, breaks at EOS and decays monotonically to zero. Overall, our results contribute to explaining two key phenomena in deep networks: (i) shallow models and simple tasks do not always exhibit EOS; and (ii) oscillations occur within top features. We present experiments to support our theory, along with examples demonstrating how these phenomena occur in nonlinear networks and how they differ from those which have benign landscape such as in DLNs.