Understanding Grokking Through A Robustness Viewpoint
This work addresses the understanding of delayed generalization in neural networks, which is an incremental advancement for researchers in machine learning theory and optimization.
The paper investigates the grokking phenomenon in neural networks, where generalization occurs long after overfitting, by analyzing it through a robustness perspective and showing that the l2 weight norm is a sufficient condition for grokking. It proposes perturbation-based methods to speed up generalization and introduces new metrics based on robustness and information theory that correlate well with grokking, potentially predicting it.
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.