Introduction

[power2022grokking] discovered that small neural networks trained on modular arithmetic tasks can exhibit "grokking" — a phenomenon where test accuracy remains at chance level long after training accuracy reaches near perfection, then suddenly jumps to high accuracy after many additional epochs of training. This delayed generalization represents a phase transition in the learning dynamics.

Subsequent work has deepened our understanding: [nanda2023progress] reverse-engineered the learned algorithm as a discrete Fourier transform with trigonometric identities, identifying three continuous training phases (memorization, circuit formation, cleanup). [liu2022omnigrok] introduced the "LU mechanism" explaining grokking through the mismatch between L-shaped training loss and U-shaped test loss as functions of weight norm, and demonstrated that grokking can be induced or suppressed across diverse data domains.

Despite these advances, a systematic mapping of the grokking phase diagram across multiple hyperparameters simultaneously has received less attention. In this work, we sweep over three key hyperparameters — weight decay, dataset fraction, and model width — to construct a complete phase diagram that classifies each training run into one of four outcomes.

Methods

Task and Data

We study modular addition: given integers $a, b \in {0, 1, \ldots, p-1}$, predict $(a + b) \bmod p$ where $p = 97$ (a standard prime used in the grokking literature). The full dataset contains $p^2 = 9,409$ input-output pairs. Each pair is unique, and labels are uniformly distributed over ${0, \ldots, p-1}$.

Model Architecture

We use a one-hidden-layer MLP with learned embeddings:

• Each input $a$ and $b$ is mapped to a 16-dimensional learned embedding
• The two embeddings are concatenated to form a 32-dimensional vector
• A linear layer maps to $h$ hidden units with ReLU activation
• A final linear layer maps to $p$ output logits

Parameter counts range from ${$ ($h = 16$) to ${}12,500$ ($h = 64$), all well under 100K.

Training

We use the AdamW optimizer [loshchilov2019adamw] with learning rate $10^{-3}$, $\beta = (0.9, 0.98)$, and cross-entropy loss. Training uses full-batch gradient descent (all training examples in each batch), following the standard grokking setup. Each run trains for up to 2,500 epochs with early stopping when both train and test accuracy exceed 99% for two consecutive evaluation points. Metrics are logged every 100 epochs.

Phase Classification

We classify each training run into one of four phases:

• Confusion: Final training accuracy $< 95$. The model fails to memorize.
• Memorization: Training accuracy $\geq 95$ but test accuracy $< 95$. Overfitting without generalization.
• Grokking: Both accuracies reach $95$, but test accuracy lags train by $> 200$ epochs. Delayed generalization.
• Comprehension: Both accuracies reach $95$ with test lagging by $\leq 200$ epochs. Fast generalization.

Hyperparameter Sweep

We sweep over:

• Weight decay $\lambda \in {0, 10^{-3}, 10^{-2}, 10^{-1}, 1.0}$ (5 values)
• Dataset fraction $f \in {0.3, 0.5, 0.7, 0.9}$ (4 values)
• Hidden dimension $h \in {16, 32, 64}$ (3 values)

This yields $5 \times 4 \times 3 = 60$ training runs, all with seed fixed to 42 for full reproducibility. The execution script also records reproducibility metadata (sweep grid, seed, package versions, and runtime) in results/metadata.json, and the validator checks full grid coverage plus phase/gap consistency for each run.

Results

Phase Diagram Structure

The phase diagram reveals clear boundaries between learning regimes. The results are presented as 2D heatmaps (weight decay $\times$ dataset fraction) for each hidden dimension.

The phase diagram shows that generalization is absent at low dataset fractions ($f = 0.3$ or $0.5$) across the entire sweep, but becomes common once $f \geq 0.7$. Within that higher-data regime, the narrowest model ($h = 16$) still struggles, while wider models display both delayed grokking and rapid comprehension. Weight decay affects which of those regimes is more likely, but no single value cleanly separates success from failure across all widths and dataset fractions.

Role of Weight Decay

Weight decay shapes the learning regime, but in this sweep it does not act as a single universal threshold:

• $\lambda = 0$: Mostly confusion or memorization, but wider models with $f \geq 0.7$ can still generalize, including two grokking runs and one rapid-comprehension run.
• $\lambda = 10^{-3}$ to $10^{-1}$: The most reliable grokking region in our sweep. These settings account for 7 of the 9 grokking runs.
• $\lambda = 1.0$: Strong regularization suppresses grokking, but does not eliminate generalization. At high dataset fraction and sufficient width, runs transition directly to comprehension instead.

Role of Dataset Fraction

Larger dataset fractions facilitate generalization. With $f = 0.3$ (30% training data), models struggle to generalize even with appropriate weight decay. At $f = 0.7$--$0.9$, grokking and comprehension are more common. This aligns with the Omnigrok finding [liu2022omnigrok] of a critical training set size below which generalization is impossible.

Role of Model Width

Wider models ($h = 64$) tend to grok or comprehend more readily than narrow models ($h = 16$) at the same weight decay and dataset fraction. This suggests that overparameterization, when combined with appropriate regularization, facilitates the transition to generalizing solutions.

Discussion

Our phase diagram confirms and refines several findings from the grokking literature:

• Dataset fraction is the clearest gate on generalization in this setup: none of the runs with $f \leq 0.5$ generalize, while most runs with $f = 0.9$ do.
• Weight decay modulates how generalization appears once enough data and width are available: intermediate values favor delayed grokking, while extreme values more often yield either confusion or fast comprehension.
• Model width modulates the phase boundaries, with wider models having broader grokking regions.
• The four-phase structure (confusion, memorization, grokking, comprehension) is robust across model widths.

Limitations. Our study uses a single arithmetic operation (addition mod 97), a single optimizer (AdamW), and a single seed. The grokking gap threshold (200 epochs) is somewhat arbitrary. Extending to multiplication, varying seeds for confidence intervals, and exploring additional optimizers would strengthen these findings.

Reproducibility. All code is deterministic on CPU with pinned seeds and dependency versions. Each run emits machine-checkable artifacts (sweep\_results.json, phase\_diagram.json, metadata.json) and passes a validator that enforces artifact presence, Cartesian grid completeness, and phase-label consistency. In our verification pass, the complete 60-run analysis finished on CPU in 383 seconds (about 6.4 minutes). No GPU, internet access, or authentication is required.

\bibliographystyle{plainnat}

References

• [power2022grokking] Power, A., Burda, Y., Edwards, H., Babuschkin, I., and Misra, V. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177, 2022.

• [nanda2023progress] Nanda, N., Chan, L., Lieberum, T., Smith, J., and Steinhardt, J. Progress measures for grokking via mechanistic interpretability. In ICLR, 2023.

• [liu2022omnigrok] Liu, Z., Michaud, E. J., and Tegmark, M. Omnigrok: Grokking beyond algorithmic data. arXiv preprint arXiv:2210.01117, 2022.

• [loshchilov2019adamw] Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In ICLR, 2019.