A Tale of Two Circuits: Grokking as Competition of Sparse and Dense Subnetworks
This work addresses the problem of understanding sudden generalization in neural networks for researchers in machine learning theory, but it is incremental as it builds on prior studies of grokking.
The paper investigates the grokking phenomenon in neural networks, where models initially overfit and then suddenly generalize perfectly after extended training, by analyzing internal network structures on the sparse parity task. It finds that grokking corresponds to the emergence of a sparse subnetwork that dominates predictions, linked to rapid norm growth in a small subset of neurons.
Grokking is a phenomenon where a model trained on an algorithmic task first overfits but, then, after a large amount of additional training, undergoes a phase transition to generalize perfectly. We empirically study the internal structure of networks undergoing grokking on the sparse parity task, and find that the grokking phase transition corresponds to the emergence of a sparse subnetwork that dominates model predictions. On an optimization level, we find that this subnetwork arises when a small subset of neurons undergoes rapid norm growth, whereas the other neurons in the network decay slowly in norm. Thus, we suggest that the grokking phase transition can be understood to emerge from competition of two largely distinct subnetworks: a dense one that dominates before the transition and generalizes poorly, and a sparse one that dominates afterwards.