🤖 AI Summary
Under scale-invariant normalization mechanisms such as RMSNorm, conventional L2 weight decay struggles to promote generalization in grokking, as it acts only along the radial direction of the weight space. This work proposes Low-Rank Decay (LRD), a spectral regularization method based on the nuclear norm that continuously compresses weight singular values through tangential components, thereby driving structural simplification of representations. For the first time, the geometric dynamics of nuclear norm regularization are introduced into grokking research, revealing that LRD achieves effective rank collapse near low-rank manifolds via a “needle-to-fan” subdifferential expansion. On modular arithmetic tasks, LRD significantly accelerates rank collapse in Query and Key matrices and substantially extends the data regime boundary at which grokking occurs, demonstrating its efficacy in facilitating delayed generalization.
📝 Abstract
Modern Transformer architectures frequently employ normalization mechanisms such as RMSNorm and Query-Key Normalization, making parts of the model approximately scale-invariant with respect to weight magnitudes. In this regime, standard Frobenius-norm weight decay acts purely along the radial direction of the weight space and cannot directly simplify the function represented by the normalized layer. We study grokking in small algorithmic tasks through this lens and propose \emph{Low-Rank Decay} (LRD), a nuclear-norm-like spectral regularizer whose subgradient -- the polar factor $UV^\top$ -- retains a tangential component even in the scale-invariant setting. This distinction has a concrete dynamical consequence: after the model memorizes the training set and task gradients vanish, L2 decay can no longer reshape the weight spectrum, whereas LRD continues to compress singular values in an $\ell_1$-like fashion. On modular arithmetic tasks, we find that LRD induces rapid effective-rank collapse in Query/Key matrices and expands the data-fraction boundary at which delayed generalization (grokking) occurs. We further provide a spectral-geometric interpretation through the ``needle-to-fan'' expansion of the nuclear-norm subdifferential near low-rank strata.