Tracing the Path to Grokking: Embeddings, Dropout, and Network Activation

TL;DR

This paper investigates the phenomenon of grokking in neural networks, introducing practical metrics like dropout variance and embedding similarity to predict delayed generalization, and analyzing neural activation patterns during this process.

Contribution

It presents new metrics such as Dropout Robustness Curve and embedding similarity measures to forecast grokking and explains neural behavior during delayed generalization.

Findings

Variance under dropout peaks during grokking

Inactive neurons decrease as models generalize

Embeddings become bimodal, indicating symmetry

Abstract

Grokking refers to delayed generalization in which the increase in test accuracy of a neural network occurs appreciably after the improvement in training accuracy This paper introduces several practical metrics including variance under dropout, robustness, embedding similarity, and sparsity measures, that can forecast grokking behavior. Specifically, the resilience of neural networks to noise during inference is estimated from a Dropout Robustness Curve (DRC) obtained from the variation of the accuracy with the dropout rate as the model transitions from memorization to generalization. The variance of the test accuracy under stochastic dropout across training checkpoints further exhibits a local maximum during the grokking. Additionally, the percentage of inactive neurons decreases during generalization, while the embeddings tend to a bimodal distribution independent of initialization that correlates with the observed cosine similarity patterns and dataset symmetries. These metrics additionally provide valuable insight into the origin and behaviour of grokking.