Muon Optimizer Accelerates Grokking
This addresses the grokking phenomenon for machine learning practitioners, offering an incremental improvement in training efficiency.
The paper tackled the problem of delayed generalization (grokking) in models by comparing optimizers, finding that the Muon optimizer significantly accelerates grokking onset, reducing the mean grokking epoch from 153.09 to 102.89 across tasks.
This paper investigates the impact of different optimizers on the grokking phenomenon, where models exhibit delayed generalization. We conducted experiments across seven numerical tasks (primarily modular arithmetic) using a modern Transformer architecture. The experimental configuration systematically varied the optimizer (Muon vs. AdamW) and the softmax activation function (standard softmax, stablemax, and sparsemax) to assess their combined effect on learning dynamics. Our empirical evaluation reveals that the Muon optimizer, characterized by its use of spectral norm constraints and second-order information, significantly accelerates the onset of grokking compared to the widely used AdamW optimizer. Specifically, Muon reduced the mean grokking epoch from 153.09 to 102.89 across all configurations, a statistically significant difference (t = 5.0175, p = 6.33e-08). This suggests that the optimizer choice plays a crucial role in facilitating the transition from memorization to generalization.