English

Weight Decay Regimes in Grokking Transformers: Cheap Online Diagnostics

Machine Learning 2026-05-21 v1 Artificial Intelligence Neural and Evolutionary Computing

Abstract

Transformers trained on modular arithmetic exhibit sharp transitions between memorization, generalization, and collapse. We show that weight decay acts as a scalar empirical control parameter for these regimes, and introduce two cheap online diagnostics, mean pairwise attention-head cosine similarity and entropy standard deviation, that track training dynamics from attention activations alone and complement loss-landscape diagnostics at lower compute cost. Across eleven experimental conditions and three model scales (0.82M to 85M parameters), the weight-decay axis separates memorization, developmental grokking, and collapse. A near-transition logistic fit localizes the memorization-to-developmental boundary at λc=0.0158\lambda_c=0.0158 (95% CI [0.0109, 0.0200], N=210); a power-law fit gives an empirical exponent ν=0.757\nu=0.757 (CI [0.725, 0.799]). Reference exponents ν=1/2\nu=1/2 and 3D Ising ν0.63\nu \approx 0.63 lie outside this empirical CI under our four-bin grid, so we report ν\nu as empirical and defer universality-class identification to denser finite-size-scaling work. A horizon-matched multi-task replication (n=280, four modular operations) preserves the weight-decay control pattern; a paired attention-head re-initialization experiment at λ=0.05\lambda=0.05 changes Phase-2 amplitude (Cohen's d=1.190d=-1.190, n=10, pt=4.5×103p_t=4.5 \times 10^{-3}), while matched weight-norm clipping does not. Three cross-architecture probes (4L MLP, 4L LSTM, and 4L Mamba; each n=70) replicate the weight-decay-controlled transition with architecture-specific λc\lambda_c values. Main diagnostic claims are scoped to modular arithmetic in small transformer attention models; the non-attention experiments are scope probes, and architecture-wide, language-model, and universality-class claims are out of scope.

Keywords

Cite

@article{arxiv.2605.20441,
  title  = {Weight Decay Regimes in Grokking Transformers: Cheap Online Diagnostics},
  author = {Lucky Verma},
  journal= {arXiv preprint arXiv:2605.20441},
  year   = {2026}
}

Comments

28 pages, 11 figures, 5 tables. Code and aggregate JSONs: https://github.com/lucky-verma/grokking-diagnostics. Per-run JSONs: https://huggingface.co/datasets/lucky-verma/grokking-diagnostics-runs. Lean 4/mathlib v4.29.0 formal checks available in the code repository