Understanding “Grokking” in Neural Networks

Neural networks often defy traditional learning paradigms. One fascinating phenomenon, explored by researchers, is grokking. In the paper, *“Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets,”* the authors reveal how neural networks can move beyond overfitting to achieve perfect generalization after extensive training suddenly.

Why Small Algorithmic Datasets?

Deep learning’s capacity for generalization defies traditional learning theory. To probe this, the researchers turned to datasets composed of equations of the form a ○ b = c, where ○ represents binary operations like modular arithmetic. These datasets are compact yet rich in structure, making them ideal for exploring how neural networks transition from memorizing data to discovering patterns.

What is Grokking?

Grokking is a phenomenon where validation accuracy starts to improve significantly well after the training accuracy has plateaued. For instance, in experiments with modular division, the authors found out that training accuracy reached near perfection after 1,000 steps, yet validation accuracy remained low until around 10⁶ steps. This delayed generalization indicates the optimization dynamics within overparameterized neural networks.

The author presents a case of generalization occurring after significant overfitting on an algorithmic dataset. The training involves the binary operation of division mod 97, utilizing 50% of the data for the training set. Each of the 97 residues is represented as a distinct symbol, akin to the representation shown in the accompanying figure. The red curves indicate training accuracy, while the green curves reflect validation accuracy. Training accuracy approaches perfection in less than 10³ optimization steps, but validation accuracy takes nearly 106 steps to achieve the same level. Generalization shows minimal evidence until around 10⁵ steps.

The time needed to achieve 99% validation accuracy grows quickly as the fraction of training data decreases.

Findings

1. Generalization Beyond Overfitting: • It is found out that the neural networks exhibited delayed validation accuracy improvements across a range of dataset sizes and binary operations. • Interestingly, smaller datasets demanded exponentially more optimization steps for the networks to generalize effectively than the larger datasets.

2. Role of Symmetry and Complexity: • The authors experimented with symmetrical operations (e.g. x + y) and found that it generalized faster than asymmetrical ones (e.g. x — y). • But there are complex operations like x³ + xy² + y mod 97 that often failed to generalize within the optimization budget in the experiment.

3. Visualization of Learned Representations:

The t-SNE projection of the output layer weights from a network trained reveals clusters of permutations. Each cluster represents a coset of the subgroup ((0, 3)(1, 4),(1, 2)(3, 4)) or one of its conjugates.

The t-SNE projection of the output layer weights from a network trained on modular addition illustrates the results of adding 8 to each element. The colors indicate the residue of each element modulo 8.

• The authors showed t-SNE projections of embeddings that revealed structured representations, such as circular topologies for modular arithmetic, to understand the networks’ internal structure of the data.

Optimization and Regularization

The study highlighted the impact of various training techniques: 〄 Weight Decay: It was found out that this regularization method greatly improved their data efficiency, thus reducing the number of training examples required for generalization. 〄 Noise-Induced Techniques: They added Gaussian noise during optimization that prompted the discovery of flatter minima, which correlate with better generalization. 〄 Hyperparameter Sensitivity: They claimed to achieve grokking, one must require precise tuning of learning rates within narrow bounds.

Mathematical Experiments and Numerical Results

The researchers explored a range of binary operations under prime modulus p = 97: • x + y mod p • x — y mod p 〄 Mixed operations, such as [ x/y if y odd, else x — y] mod p.

They observed that: ⦿ Smaller training datasets required exponentially longer optimization times. ⦿ Validation loss curves often displayed a “double descent,” that is, a phenomenon where loss initially increases before improving greatly.

Discussion and Future Directions

Their preliminary evidence suggests that grokking may occur when parameters converge to flatter regions of the loss landscape. While this study focused on algorithmic datasets, future work could explore whether grokking generalizes to real-world data. So, for investigating the role of neural network architecture for data efficiency, grokking remains an open question

The concept of grokking offers us insight into how neural networks generalize. This research improves our theoretical knowledge and establishes a foundation for practical progress in machine learning. The results encourage us to investigate the circumstances that allow neural networks to move beyond simple memorization and attain true pattern recognition.