🤖 AI Summary
This work investigates the delayed generalization—commonly referred to as grokking—in neural networks trained on algorithmic tasks, which often arises from over-memorization of training data. From a geometric perspective, the authors identify radial expansion in activation space as a key contributor to this phenomenon. They propose a radial-angular decomposition framework and introduce a simple L2-norm activation constraint with only one hyperparameter, softly confining hidden representations to a hypersphere of radius √d. This mechanism suppresses radial gradients and induces anisotropic regularization, thereby steering optimization toward angular directions. Consequently, the network more rapidly discovers structured low-dimensional circuits and converges to flatter minima. Empirically, the method achieves up to a 6× acceleration in grokking on modular arithmetic tasks and reduces training steps by 50% on three-digit addition with a 10M-parameter nanoGPT model.
📝 Abstract
Why do neural networks memorize algorithmic training data long before they generalize? We present a geometric case study demonstrating that, on tasks where generalization requires discovering structured low-dimensional circuits, the memorization-generalization delay is driven by radial inflation of hidden representations under cross-entropy optimization. We formalize a radial-angular decomposition of activation-space dynamics and derive three testable propositions: (i) that penalizing radial inflation induces anisotropic, data-dependent weight regularization; (ii) that it suppresses radial gradient energy below the isotropic random baseline, forcing predominantly angular updates; and (iii) that it biases convergence toward flatter minima. To empirically validate these propositions, we study a single-hyperparameter norm penalty that softly constrains activations to a sqrt(d)-radius hypersphere. On modular arithmetic, this penalty accelerates grokking up to 6x across MLPs and Transformers, and halves training steps for a 10M-parameter nanoGPT on 3-digit addition.