Introduction

Grokking—the phenomenon where neural networks first memorize training data, then suddenly generalize after extended training—has attracted significant attention since its discovery by [power2022grokking]. Understanding when and why this generalization transition occurs remains an active area of research.

Prior work has connected grokking to weight norm dynamics [liu2022omnigrok], representation learning phase transitions [nanda2023grokking], and the lazy-to-rich training regime transition [lyu2024grokking]. However, the relationship between per-layer gradient norm dynamics and the onset of generalization has received less direct attention.

We hypothesize that gradient norm phase transitions—specifically, the peak of gradient $L_2$ norms during training—serve as an early indicator of the memorization-to-generalization transition. If gradient norms peak and begin declining before test accuracy improves, they could function as a "leading indicator" for generalization, useful for early stopping decisions and training diagnostics.

Methods

Tasks and Models

We study two tasks:

- **Modular addition** (mod 97): Given one-hot encoded <span class="katex"><span class="katex-mathml"><math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><mo stretchy="false">(</mo><mi>a</mi><mo separator="true">,</mo><mi>b</mi><mo stretchy="false">)</mo></mrow><annotation encoding="application/x-tex">(a, b)</annotation></semantics></math></span><span class="katex-html" aria-hidden="true"><span class="base"><span class="strut" style="height:1em;vertical-align:-0.25em;"></span><span class="mopen">(</span><span class="mord mathnormal">a</span><span class="mpunct">,</span><span class="mspace" style="margin-right:0.1667em;"></span><span class="mord mathnormal">b</span><span class="mclose">)</span></span></span></span>, predict <span class="katex"><span class="katex-mathml"><math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><mo stretchy="false">(</mo><mi>a</mi><mo>+</mo><mi>b</mi><mo stretchy="false">)</mo><mtext> </mtext><mo lspace="0.22em" rspace="0.22em"><mrow><mi mathvariant="normal">m</mi><mi mathvariant="normal">o</mi><mi mathvariant="normal">d</mi></mrow></mo><mtext> </mtext><mn>97</mn></mrow><annotation encoding="application/x-tex">(a + b) \bmod 97</annotation></semantics></math></span><span class="katex-html" aria-hidden="true"><span class="base"><span class="strut" style="height:1em;vertical-align:-0.25em;"></span><span class="mopen">(</span><span class="mord mathnormal">a</span><span class="mspace" style="margin-right:0.2222em;"></span><span class="mbin">+</span><span class="mspace" style="margin-right:0.2222em;"></span></span><span class="base"><span class="strut" style="height:1em;vertical-align:-0.25em;"></span><span class="mord mathnormal">b</span><span class="mclose">)</span><span class="mspace" style="margin-right:0.0556em;"></span><span class="mspace" style="margin-right:0.2222em;"></span><span class="mbin"><span class="mord"><span class="mord mathrm">mod</span></span></span><span class="mspace" style="margin-right:0.0556em;"></span><span class="mspace" style="margin-right:0.2222em;"></span></span><span class="base"><span class="strut" style="height:0.6444em;"></span><span class="mord">97</span></span></span></span>. This is a standard grokking benchmark [power2022grokking].
- **Polynomial regression**: Predict <span class="katex"><span class="katex-mathml"><math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><mi>y</mi><mo>=</mo><mi>sin</mi><mo>⁡</mo><mo stretchy="false">(</mo><mi>x</mi><mo stretchy="false">)</mo><mo>+</mo><mn>0.3</mn><mi>sin</mi><mo>⁡</mo><mo stretchy="false">(</mo><mn>3</mn><mi>x</mi><mo stretchy="false">)</mo></mrow><annotation encoding="application/x-tex">y = \sin(x) + 0.3\sin(3x)</annotation></semantics></math></span><span class="katex-html" aria-hidden="true"><span class="base"><span class="strut" style="height:0.625em;vertical-align:-0.1944em;"></span><span class="mord mathnormal" style="margin-right:0.0359em;">y</span><span class="mspace" style="margin-right:0.2778em;"></span><span class="mrel">=</span><span class="mspace" style="margin-right:0.2778em;"></span></span><span class="base"><span class="strut" style="height:1em;vertical-align:-0.25em;"></span><span class="mop">sin</span><span class="mopen">(</span><span class="mord mathnormal">x</span><span class="mclose">)</span><span class="mspace" style="margin-right:0.2222em;"></span><span class="mbin">+</span><span class="mspace" style="margin-right:0.2222em;"></span></span><span class="base"><span class="strut" style="height:1em;vertical-align:-0.25em;"></span><span class="mord">0.3</span><span class="mspace" style="margin-right:0.1667em;"></span><span class="mop">sin</span><span class="mopen">(</span><span class="mord">3</span><span class="mord mathnormal">x</span><span class="mclose">)</span></span></span></span> from <span class="katex"><span class="katex-mathml"><math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><mi>x</mi><mo>∈</mo><mo stretchy="false">[</mo><mo>−</mo><mn>3</mn><mo separator="true">,</mo><mn>3</mn><mo stretchy="false">]</mo></mrow><annotation encoding="application/x-tex">x \in [-3, 3]</annotation></semantics></math></span><span class="katex-html" aria-hidden="true"><span class="base"><span class="strut" style="height:0.5782em;vertical-align:-0.0391em;"></span><span class="mord mathnormal">x</span><span class="mspace" style="margin-right:0.2778em;"></span><span class="mrel">∈</span><span class="mspace" style="margin-right:0.2778em;"></span></span><span class="base"><span class="strut" style="height:1em;vertical-align:-0.25em;"></span><span class="mopen">[</span><span class="mord">−</span><span class="mord">3</span><span class="mpunct">,</span><span class="mspace" style="margin-right:0.1667em;"></span><span class="mord">3</span><span class="mclose">]</span></span></span></span>. This serves as a smooth-learning control task.

Both tasks use a 2-layer MLP: input $\to$ Linear(hidden=64) $\to$ ReLU $\to$ Linear(output) $\to$ output.

Training Configuration

We train with AdamW (lr=$10^{-3}$, weight decay=$0.1$) for 3000 epochs on three dataset fractions: 50%, 70%, and 90%. All runs use seed 42 for the primary sweep and seeds {42, 123, 7} for variance analysis. For reproducibility, we enable deterministic PyTorch algorithms and log runtime/platform/library metadata into the output JSON. This yields $2 \times 3 = 6$ primary training runs plus $3 \times 3 = 9$ variance runs on modular addition.

Metrics Tracked

At every epoch, we record:

- Per-layer gradient <span class="katex"><span class="katex-mathml"><math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><msub><mi>L</mi><mn>2</mn></msub></mrow><annotation encoding="application/x-tex">L_2</annotation></semantics></math></span><span class="katex-html" aria-hidden="true"><span class="base"><span class="strut" style="height:0.8333em;vertical-align:-0.15em;"></span><span class="mord"><span class="mord mathnormal">L</span><span class="msupsub"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height:0.3011em;"><span style="top:-2.55em;margin-left:0em;margin-right:0.05em;"><span class="pstrut" style="height:2.7em;"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight">2</span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height:0.15em;"><span></span></span></span></span></span></span></span></span></span> norms (after backward pass, before optimizer step)
- Per-layer weight <span class="katex"><span class="katex-mathml"><math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><msub><mi>L</mi><mn>2</mn></msub></mrow><annotation encoding="application/x-tex">L_2</annotation></semantics></math></span><span class="katex-html" aria-hidden="true"><span class="base"><span class="strut" style="height:0.8333em;vertical-align:-0.15em;"></span><span class="mord"><span class="mord mathnormal">L</span><span class="msupsub"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height:0.3011em;"><span style="top:-2.55em;margin-left:0em;margin-right:0.05em;"><span class="pstrut" style="height:2.7em;"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight">2</span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height:0.15em;"><span></span></span></span></span></span></span></span></span></span> norms
- Train/test loss and train/test accuracy (or <span class="katex"><span class="katex-mathml"><math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><msup><mi>R</mi><mn>2</mn></msup></mrow><annotation encoding="application/x-tex">R^2</annotation></semantics></math></span><span class="katex-html" aria-hidden="true"><span class="base"><span class="strut" style="height:0.8141em;"></span><span class="mord"><span class="mord mathnormal" style="margin-right:0.0077em;">R</span><span class="msupsub"><span class="vlist-t"><span class="vlist-r"><span class="vlist" style="height:0.8141em;"><span style="top:-3.063em;margin-right:0.05em;"><span class="pstrut" style="height:2.7em;"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight">2</span></span></span></span></span></span></span></span></span></span></span>)

Phase Transition Detection

We detect gradient norm transitions using the peak epoch: the epoch at which the smoothed (Savitzky-Golay filter, window=51) combined gradient norm reaches its maximum, skipping the initial 2% of training to avoid transient effects.

Test metric transitions are detected as the epoch of steepest increase in the smoothed test accuracy (or $R^2$).

The lag is defined as: $\text{lag} = \text{epoch}$\text{metric transition} - \text{epoch}\text{gradient peak}. Positive lag means gradient norms transition first.

We additionally compute cross-correlation between the derivative of the gradient norm signal and the derivative of the test metric to quantify temporal coupling, with deterministic tie-breaking toward zero lag when correlations are equal.

Results

Modular Addition: Gradient Norms Lead Generalization

Phase transition epochs and lag for all experimental runs. Positive lag indicates the gradient norm peak precedes the test metric transition.

Table summarizes the key findings. For all three modular addition configurations, the gradient norm peak precedes the test accuracy transition:

- At 50% data fraction, gradient norms peak at epoch 562, while test accuracy remains low by epoch 3000 (6.6%). The steepest test-accuracy increase is still detected later, at epoch 1261, so the gradient peak leads by 699 epochs. This suggests the internal dynamics shift well before strong held-out generalization is visible.
- At 70%, the peak-to-generalization lag is 192 epochs, and the model reaches 76.2% test accuracy.
- At 90%, the lag is 12 epochs—still positive but small, consistent with the faster generalization observed with more training data.

Regression: No Delayed Generalization, No Lag

The regression task shows no grokking. Both gradient norms and test $R^2$ transition immediately (test metric at epoch 0, gradient peak at epoch 60). The negative lag ($-60$) reflects the absence of a memorization phase: the model generalizes from the start, and gradient norms simply follow an initial rise-and-decay pattern with no predictive power.

Correlation Structure

The Pearson correlation between gradient norm and test metric is strongly negative ($r \in [-0.93, -0.66]$, $p \approx 0$ for all runs), confirming that gradient norm decline is temporally anti-correlated with performance improvement. The correlation is strongest in tasks with more pronounced grokking.

Multi-Seed Variance Analysis

To assess robustness, we repeat the modular addition experiments across 3 random seeds (42, 123, 7). Table shows that the gradient-leading-metric pattern is consistent across seeds, with the lag always positive.

Multi-seed lag statistics for modular addition (3 seeds). The gradient norm peak consistently leads the test accuracy transition.

The variance is highest at 50% data fraction, where grokking dynamics are most sensitive to initialization. At 70% and 90%, the lag is more consistent across seeds (CV of 14% and 62% respectively).

Discussion

Our results support the hypothesis that gradient norm phase transitions—specifically, the peak of gradient $L_2$ norms—precede generalization transitions in grokking-prone settings. The gradient norm peak marks the point where the network's optimization landscape shifts from building memorization circuits (high gradient activity) to consolidating generalization circuits (declining gradients as weight decay regularization takes effect).

The monotonic relationship between data fraction and lag is noteworthy: less training data produces a larger lag ($642 \pm 197$ epochs at 50% vs.\ $54 \pm 37$ epochs at 90%). This aligns with the theoretical picture where weight decay must work longer to overcome the memorization solution when data is scarce [liu2022omnigrok]. The multi-seed analysis confirms this pattern is robust to initialization.

Limitations

- We study only 2-layer MLPs; deeper architectures may show different layer-wise dynamics.
- Only 3 seeds are used for variance analysis; larger ensembles would provide tighter confidence intervals.
- The transition detection uses a smoothing heuristic; more principled change-point detection methods could improve robustness.
- Only two tasks are studied; broader task families (e.g., group operations beyond addition, image classification) would strengthen generalizability claims.

Conclusion

Gradient $L_2$ norm peaks serve as an early warning signal for the memorization-to-generalization transition in grokking-prone tasks, preceding test accuracy improvements by 12--699 epochs in the primary runs and by 54--642 epochs on average in the 3-seed modular-addition sweeps. This finding has practical implications for training diagnostics: monitoring gradient norm trajectories could allow practitioners to predict whether and when a model will generalize, without waiting for test performance to improve.

\bibliographystyle{plainnat}

References

• [power2022grokking] A. Power, Y. Burda, H. Edwards, I. Babuschkin, and V. Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. In ICLR 2022 MATH-AI Workshop, 2022.

• [liu2022omnigrok] Z. Liu, E. J. Michaud, and M. Tegmark. Omnigrok: Grokking beyond algorithmic data. In ICLR, 2023.

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

• [lyu2024grokking] K. Lyu, J. Jin, Z. Li, S. S. Du, J. D. Lee, and W. Hu. Grokking as the transition from lazy to rich training dynamics. In ICLR, 2024.