Radial Suppression Accelerates Algorithmic Generalization: A Geometric Analysis of Delayed Generalization
For researchers studying grokking and generalization in neural networks, this work provides a geometric explanation and a simple method to accelerate generalization, though it is an incremental improvement over existing regularization techniques.
The paper identifies radial inflation of hidden representations as the cause of delayed generalization (grokking) in neural networks, and shows that a norm penalty constraining activations to a hypersphere accelerates generalization by up to 6x on modular arithmetic and halves training steps for a 10M-parameter nanoGPT on 3-digit addition.
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.