Compute-Optimal Scaling Laws for the Generalization Phase Transition in Grokking
Anish Kataria
Abstract
We derive compute-optimal scaling laws for the generalization phase transition known as grokking. A 384-configuration sweep of two-layer MLPs on modular arithmetic yields $T_{\mathrm{grok}} \propto H^{-0.27} D^{-2.04} \eta^{-0.50} \lambda^{-0.64}$ ($R^2 = 0.73$; $0.82$ with interactions), where $H$ is width (exponent from three width levels), $D$ dataset fraction, $\eta$ learning rate, and $\lambda$ weight decay. Dataset fraction dominates: doubling $D$ cuts grokking time by ~4x. A phase diagram in $(\eta, \lambda)$ space reveals a sharp boundary separating grokking from non-grokking regimes. Compute-optimal analysis shows wider models grok in fewer steps but at higher FLOP cost, mirroring Chinchilla-style trade-offs. The law predicts generalization onset from hyperparameters alone, before training begins, enabling practitioners to set compute budgets without pilot runs.
Chat is not available.
Successful Page Load