Simplicity Bias of SGD via Sharpness Minimization

Khashayar Gatmiry, Zhiyuan Li, Luana Ruiz, Sashank J. Reddi, Stefanie Jegelka

Primary Area: learning theory

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: Sharpness minimization, Implicit bias, SGD, Simplicity Bias, trace of Hessian regularizer

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Abstract:

The remarkable generalization ability of neural networks is usually attributed to the implicit bias of SGD, which often yields models with lower complexity using simpler (e.g. linear) and low-rank features (Huh et al., 2021). Recent works have provided empirical and theoretical evidence for the bias of particular variants of SGD (such as label noise SGD) towards flatter regions of the loss landscape.
Despite the folklore intuition that flat solutions are 'simple', the connection with the simplicity of the final trained model (e.g. low rank) is not well understood.
In this work, we take a step towards bridging this gap by studying the simplicity structure that arises from minimizers of the sharpness for a class of two-layer neural networks.
We show that, for any high dimensional training data and certain activations, with small enough step size, label noise SGD always converges to a network that replicates a single linear feature across all neurons; thereby implying a simple rank one feature matrix. To obtain this result, our main technical contribution is to show that label noise SGD always minimizes the sharpness on the manifold of models with zero loss for two-layer networks.
Along the way, we discover a novel property --- a local geodesic convexity --- of the trace of Hessian of the loss at approximate stationary points on the manifold of zero loss, which links sharpness to the geometry of the manifold. This tool may be of independent interest.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 8659