Understanding Grokking Through A Robustness Viewpoint

TL;DR

The paper tackles the grokking phenomenon, where neural networks generalize long after overfitting, by framing it through network robustness and the decay of the $l_2$ weight norm as a sufficient condition for grokking. It introduces perturbation-based training to speed generalization and reveals that learning fundamental group properties, such as commutativity in modulo addition, can be essential for grokking on certain tasks. To better predict and understand grokking, the authors develop matrix-information-theoretic metrics (PMI and PE) and their robustness-based variants (MID, ED) that correlate strongly with test performance. The work also demonstrates that these strategies can alter representations and provides a theoretical link between robustness, phase-transition-like generalization, and practical acceleration of grokking, with broad implications for diagnostic tools and training regimes.

Abstract

Recently, an interesting phenomenon called grokking has gained much attention, where generalization occurs long after the models have initially overfitted the training data. We try to understand this seemingly strange phenomenon through the robustness of the neural network. From a robustness perspective, we show that the popular $l_2$ weight norm (metric) of the neural network is actually a sufficient condition for grokking. Based on the previous observations, we propose perturbation-based methods to speed up the generalization process. In addition, we examine the standard training process on the modulo addition dataset and find that it hardly learns other basic group operations before grokking, for example, the commutative law. Interestingly, the speed-up of generalization when using our proposed method can be explained by learning the commutative law, a necessary condition when the model groks on the test dataset. We also empirically find that $l_2$ norm correlates with grokking on the test data not in a timely way, we propose new metrics based on robustness and information theory and find that our new metrics correlate well with the grokking phenomenon and may be used to predict grokking.

Figures (13)

• Figure 1: Typical grokking cases on MNIST and Modulo Addition Dataset: Grokking phenomena refer to sudden and unforeseen enhancements of test accuracy far beyond the point when training accuracy reaches 100%.
• Figure 2: Robustness and sharpness on MNIST and Modulo Addition training dataset, we use Algorithmic to indicate Modulo Addition dataset to save the space for image caption. (Set the deviation of Gaussian perturbation as $\sigma= 0.04$). We plot the perturbation error as a robustness indicator in (a) and (b), a smaller perturbation error means more robust.
• Figure 3: The predicted accuracy matches the real test accuracy.
• Figure 4: Our perturbed training strategy speeds up generalization ("degrokking").
• Figure 5: Abelian test shows that perturb training will help model comprehend commutative rule.

Theorems & Definitions (29)

• Definition 2.1: $\alpha$-order matrix entropy
• Definition 2.2: Matrix mutual information
• Lemma 3.1
• Theorem 3.2
• Corollary 3.3
• Corollary 3.4
• proof
• Definition 5.1: Perturb Mutual Information
• Definition 5.2: Perturb Entropy
• Definition 5.3