Biologically Inspired Mechanisms Promote Grokking in Neural Networks

The phenomenon of "grokking" in neural networks, where a model transitions from memorizing training data to genuinely generalizing after extended training, is often accompanied by significant internal representation changes. This research explores whether incorporating mechanisms inspired by biological brains can actively encourage this transition. The study focuses on regulating hidden-layer computation at various levels, including neuronal activity, response, and effective connectivity. The researchers augmented a multilayer perceptron with several biologically inspired features, such as input gating, structural plasticity, gain modulation, threshold modulation, homeostasis, lateral inhibition, and activation decorrelation. They systematically evaluated these mechanisms through ablations on established grokking benchmarks like sparse parity and noisy XOR classification. The findings indicate that these mechanisms contribute unequally to generalization. Homeostasis emerged as the most potent and consistent benefit, while structural sparsification was identified as the second major contributing mechanism. The study supports the principle that explicit regulation of neuron utilization and effective connectivity can foster the emergence of generalizable internal computation, suggesting potential applications for accelerating robust generalization in larger models, including Large Language Models.