Grokking at the Edge of Linear Separability

TL;DR

This paper explores the grokking phenomenon in a simple binary classification task, revealing it occurs near critical points where training dynamics linger, and draws parallels to physical critical phenomena.

Contribution

It provides the first analytical and empirical explanation of grokking in a minimal model, linking it to critical points and implicit bias of gradient descent.

Findings

Grokking occurs near the interpolation threshold in simple models.

Flat directions in the loss landscape cause prolonged training dynamics.

Grokking is related to critical phenomena observed in physical systems.

Abstract

We investigate the phenomenon of grokking -- delayed generalization accompanied by non-monotonic test loss behavior -- in a simple binary logistic classification task, for which "memorizing" and "generalizing" solutions can be strictly defined. Surprisingly, we find that grokking arises naturally even in this minimal model when the parameters of the problem are close to a critical point, and provide both empirical and analytical insights into its mechanism. Concretely, by appealing to the implicit bias of gradient descent, we show that logistic regression can exhibit grokking when the training dataset is nearly linearly separable from the origin and there is strong noise in the perpendicular directions. The underlying reason is that near the critical point, "flat" directions in the loss landscape with nearly zero gradient cause training dynamics to linger for arbitrarily long times near quasi-stable solutions before eventually reaching the global minimum. Finally, we highlight similarities between our findings and the recent literature, strengthening the conjecture that grokking generally occurs in proximity to the interpolation threshold, reminiscent of critical phenomena often observed in physical systems.