Early-Warning Signals of Grokking via Loss-Landscape Geometry
This provides a robust, architecture-agnostic early-warning signal for delayed generalization in transformers, addressing a key challenge in understanding and controlling learning dynamics in machine learning.
The study tackled the problem of predicting grokking—the abrupt transition from memorization to generalization—in sequence-learning tasks like SCAN and Dyck-1, finding that the commutator defect, a curvature measure, rises before generalization with lead times following a superlinear power law (alpha ~1.18 for SCAN, ~1.13 for Dyck), and causal interventions showed amplifying non-commutativity accelerates grokking by roughly 32% on SCAN and 50% on Dyck.
Grokking -- the abrupt transition from memorization to generalization after prolonged training -- has been linked to confinement on low-dimensional execution manifolds in modular arithmetic. Whether this mechanism extends beyond arithmetic remains open. We study two sequence-learning benchmarks: SCAN compositional generalization and Dyck-1 depth prediction. Across both tasks and a wide range of learning rates, the commutator defect -- a curvature measure derived from non-commuting gradient updates -- rises well before generalization, with lead times following a superlinear power law (alpha approximately 1.18 for SCAN, approximately 1.13 for Dyck), consistent with prior results on modular arithmetic. Weight-space PCA reveals that spectral concentration is not a universal precursor; the commutator defect is. Causal interventions demonstrate a mechanistic role: amplifying non-commutativity accelerates grokking (roughly 32% on SCAN, roughly 50% on Dyck), while suppressing orthogonal gradient flow delays or prevents it. The three task families form a spectrum of causal sensitivity -- modular arithmetic is rigid, Dyck is responsive, SCAN is intermediate -- yet suppression delays or prevents grokking in all cases, establishing necessity as a universal finding. These results identify the commutator defect as a robust, architecture-agnostic, causally implicated early-warning signal for delayed generalization in transformers.