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

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 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 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. Finally, insight into the origin and behaviour of grokking is obtained by applying the above metrics to modular addition and MNIST digit analysis in both multi-layer perceptron (MLP) and long short-term memory (LSTM) networks.