Flat Loss, Evolving Features: Grokking in Modular Addition

Akash Kumar; Karthik Sankaranarayanan

cs.LGOct 15, 2026

Flat Loss, Evolving Features: Grokking in Modular Addition

Authors: Akash Kumar, Karthik Sankaranarayanan

Organizations: University of California, San Diego · Microsoft AI

Abstract

Can a network keep learning while its training cross-entropy stays nearly flat at a positive level? We show that it can. On modular addition, a standard testbed for grokking, ReLU networks trained with supported-Polar hidden updates, normalized output-layer updates and weight decay can keep 100% training accuracy while their hidden features change and held-out accuracy rises from zero to above 90%, at relatively flat training cross-entropy. Correct classification need not minimize cross-entropy, which still depends on prediction confidence. We prove this coexistence for one-hidden-layer networks in two finite-window theorems from constructed initializations: one gives a constant gain in the Fourier fraction of the hidden features, and the other holds at arbitrarily low positive loss, where a readout refitted on the training data generalizes before the network. Experiments from random initialization show the same coexistence on a window fixed in advance, and in deeper networks held-out accuracy also rises over post-hoc windows with a net increase in training cross-entropy. With a last-layer rate set per training fraction and head decay scaled with the prime, deeper networks at p=47 and 61 meet a criterion on post-hoc windows, fixed in advance, in nearly all fresh runs. A one-block transformer trained with the same updates and a slowly decaying output gain keeps every training pair correct at every logged checkpoint while its training cross-entropy rises and its held-out accuracy approaches one; with two blocks this holds through the held-out rise, after which training pairs begin to flip. A plateau in training loss need not be a plateau in learning. Code: https://github.com/akashkumar-d/flat-loss-evolving-features.

Figures & tables

Figure 1: Coexistence on a fixed window (1HL ReLU network,

p=31

, width 512, iid initialization, head decay

\lambda_{V}=0.2

). All 20 runs, each with a fresh initialization and training split (thin), and the median (thick). From the left: training and held-out CE; the training CE on its own scale; training accuracy (dashed) and the held-out accuracy of the network and of the train-only ridge refit; the Fourier fraction

A_{H}

. Gray: the locked window 4.5k–12k; dashed line: median per-run loss peak. Loss curves show cycle means. On the window, every run keeps perfect training accuracy, has raw per-update loss variation

6.33

–

7.81\%

and reaches

90\%

held-out accuracy. In every run, a few training pairs are misclassified again for a while after the first fit, long before the window (Appendix B.2 ).

Figure 2: One hidden layer across primes (width 512, 10 runs per prime;

\lambda_{W}=0.01,0.015,0.02,0.02,0.02

p=23,31,47,61,101

, chosen on development seeds, for

p=101

p=97

). Rows: cycle-mean training CE divided by

\log p

; median training accuracy (dashed), held-out accuracy of the network (runs and median) and of the train-only refit (median);

A_{H}

. Gray: from the median start to the median end of the runs’ post-hoc windows (Appendix A.3 ). The primary endpoint of Section 5.2 is met by 9, 4, 0, 0 and 0 runs; at

p\geq 47

some training pairs stay misclassified inside the window. Runs stopped early by the stopping rule keep their last value. Dotted (

p=31

): medians with a hidden bias (Appendix C.6 ).

Figure 3: Two to five hidden layers at

p=31

: every run (thin) and the median (thick). Width 256. Depth 2 uses a fixed rate set on development runs; depths 3–5 tune the last-layer rate on a preliminary run (Appendix C.1 ); every depth shows the same 40 evaluation seeds. The primary endpoint is met by 34, 40, 39 and 38 runs; the six depth-2 misses fail the share threshold, and one of them also ends with held-out accuracy below

0.99

. Rows as in Figure 2 , with the last-hidden-layer

A_{H}

(solid) and the median first-hidden-layer

A_{H}

(dash-dot). Gray: from the median start to the median end of the runs’ post-hoc windows (Appendix A.3 ), which contain the rise in loss after the knee. One hidden layer at

p=31

is the second column of Figure 2 . Dotted (depths 2, 4): medians with hidden biases.

Figure 4: Depth 3 at

p=47

across training fractions: every run (thin) and the median (thick). Width 256, 10 fresh runs per fraction, last-layer multiplier

r=0.45

0.30

0.18

and

0.08

chosen per fraction on development runs, head decay

\lambda_{V}=0.2\cdot 31/p

(Appendix C.4 ). Rows and gray bands as in Figure 3 . The criterion of Appendix C.4 , fixed before the runs, is met by 10, 10, 10 and 8 runs; the two misses end at held-out accuracy

0.959

and

0.971

when they reach the update budget. No training pair is misclassified at any log inside any window. Each column ends no later than its shortest run.

Figure 5: Transformers with one and two blocks: every run (thin) and the median (thick). One block at

p=89

(training fraction

0.25

\lambda_{V}=0.33

) and at

p=113

(

0.2

\lambda_{V}=0.275

), with the output gain decaying as

g_{t}\approx 1+e^{-t/5000}

, and two blocks at

p=113

(

0.2

\lambda_{V}=0.275

g_{t}\approx 1+e^{-t/10000}

); 10 runs per column (seeds

i=11

–

20

). Rows as in Figure 2 , with the training CE averaged over 201 updates (each column on its own scale) and

A_{H}

of the normalized head input

\hat{r}

. The accuracy row adds the two medians used by the pass rule of Section 5.3 : held-out accuracy on the no-mirror pairs (dashed) and the Fourier-filtered refit, which replaces the refit of Figure 2 . The refit, not drawn, trails the network at its

50\%

crossing in all 30 runs, by a median of 0.7k, 1.2k and 0.5k updates in the three columns. Gray: from the median start to the median end of the post-hoc windows on all held-out pairs (Appendix D.2 ). Each column ends at its shortest run. One block passes in 10 of 10 runs at both primes; its runs stop at 8.2k–10.2k updates (

p=89

) and 9.3k–14k (

p=113

). Two blocks pass in none. Through the held-out rise they behave like one block: the CE rises and every training pair is correct at every 20-update log. The first training errors are logged at 17.8k–22.0k updates, 1.4k to 5.7k updates after held-out accuracy first reaches

0.99

; held-out accuracy then falls back and jitters, and 6 of the 10 runs end below

0.99

(Appendix D.3 , Figure 15 ).

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Study	$c/\eta_{V}/\lambda_{W}$	Window	$p$	Runs	Flat	Rec.	First-fit	Med. $v$	Med. $\Delta A_{H}$	Med. min CE
Headline (fresh)	24/.03/.015	pre $[4.5,12]$ k	31	20	20	20	0	7.06%	.0965	1.48
Per-prime (fresh)	24/.01/.005	pre $[13,21]$ k	17	20	16	20	15	6.3%	.039	$1.8\cdot 10^{-4}$
	48/.03/.005	pre $[15,21]$ k	23	20	20	20	0	5.4%	.041	$5.1\cdot 10^{-3}$
	96/.03/.005	pre $[18,24]$ k	29	20	20	20	0	6.5%	.043	$3.9\cdot 10^{-2}$
Replays	24/.03/.01	per run	23	5	5	5	–	2.4%	.038	0.26
	24/.03/.015	per run	31	5	5	5	–	2.7%	.049	1.5

Appendix

Table 1: One-hidden-layer studies (ReLU rows: width 512,

80\%

random split,

\eta_{W}=0.03

\lambda_{V}=0.2

). Window: pre = fixed on development runs before the fresh pairs; per run = persistent held-out

10\%\!\to\!90\%

crossings (descriptive). Flat:

v\leq 10\%

. Rec.: held-out

90\%

within the horizon. First-fit: the strict joint endpoint of Appendix A.3 ; in particular, every training margin must stay above

10^{-10}

from

T_{\rm fit}

through recovery and the window end. –: not computed. The fixed-recipe sweep is in Appendix B.4 .

Figure 6: Pre-registered

p=23

cohorts (20 fresh pairs each; window

I=[6500,16500]

chosen on four development pairs and locked before these runs). (a) Cycle-mean training CE relative to its value at

t_{0}

; dots mark each run’s first

90\%

(top) or

50\%

(bottom) held-out crossing on

I

. (b) Median training accuracy (dashed) and held-out accuracy of the network and of the train-only refit (500-update grid). (c)

A_{H}

. Thin curves are runs and thick curves medians. The joint endpoint, which limits cycle-mean drawdown to

0.5\%

before the dot and requires the loss at the dot to be no smaller than the loss at

t_{0}

, holds in 10 of 20 runs (top) and 20 of 20 (bottom); 18 were required.

Figure 7: Locked and pre-registered one-hidden-layer cohorts (20 fresh pairs per column). (a)

p=23

with the headline hyperparameters: the second pre-registered cohort, window

[6500,16500]

. (b)–(d) The per-prime recipes at

p=17,23,29

, with windows

[13000,21000]

[15000,21000]

and

[18000,24000]

. Rows: cycle-mean training CE divided by

\log p

(each column on its own scale); median training accuracy (dashed), held-out accuracy of the network (runs and median) and median train-only refit;

A_{H}

. Gray: from the median start to the median end of the runs’ post-hoc windows. The headline runs at

p=31

are in Figure 1 .

Figure 8: Theorem-inspired settings at

p=101

from iid initialization (width 512; all runs in the two displayed arms). Rows: head steps

\eta_{V}=3\cdot 10^{-4}

(10 runs; all criteria of Appendix B.6 hold in each) and the primary

\eta_{V}=10^{-3}

(10 runs; the training loss varies by

1.34

–

1.37\%

over

[1000,40000]

, above the

1\%

criterion). Columns: training cross-entropy loss divided by

\log p

(the axis spans

2\%

\log p

); median training accuracy (dashed), held-out accuracy of the network (runs and median) and median held-out accuracy of the theorem’s train-only refit;

A_{H}

. Gray: from the median first observed checkpoint at which the refit’s held-out accuracy reaches

0.1

to the median first at which it reaches

0.9

; these medians need not define any individual run’s interval.

Figure 9: Results for all 164 evaluated deep-network runs. Dots are runs and bars medians. Left of each divider: the primary cohorts (

p=31

, width 256,

D=3,4,5

); right:

D=4

with

p=23

D=4

with width 512, and

D=3

with Frobenius-normalized hidden updates (NGD). (a) Share of the

10\%

-to-

90\%

held-out rise inside the loss-only band

I_{C}

, with primary-endpoint pass counts above; (b) last-hidden-layer Fourier-fraction gain

\Delta A_{H}

I_{C}

; (c) held-out accuracy at the saved horizon. Dashed lines mark the primary thresholds in (a)–(c); the fourth primary criterion, positive logged training margins in

I_{C}

, holds in all 164 runs. Rates were selected by rule ( 26 ) .

Setting	$(D,M,p)$	Indices $i$	$r_{0}$	$x_{*}$	Horizon	Primary	$J_{1}/J_{\rm pers}$	$P_{10}/P_{3}$
Primary	$(3,256,31)$	21–40 / 1–20	.50	3.85	80k	40/40	0/17	40/2
Primary	$(4,256,31)$	21–40 / 1–20	.35	3.85	100k	39/40	0/13	39/0
Primary	$(5,256,31)$	21–40 / 1–20	.35	3.85	80k	38/40	0/19	39/0
Modulus	$(4,256,23)$	41–60	.35	3.85	120k	0/20	0/0	0/0
Width	$(4,512,31)$	61–64	.40	4.50	100k	0/4	0/1	4/1
NGD	$(3,256,31)$	1–20	.50	3.85	80k	0/20	0/0	0/0

Appendix

Table 2: Deep-network settings and outcomes, retaining all evaluated runs. Rate-selection runs stop at

t_{10}^{\rm sel}

, except those also used for fixed-rate comparisons, which continue to the full horizon. Each evaluated run restarts from the same initialization and split used to select its rate. Primary: the per-run primary endpoint.

J_{1}/J_{\rm pers}

: first-fit and persistent-fit strict joints;

P_{10}/P_{3}

: raw-CE range at most

10\%

3\%

from the first

10\%

to the first

90\%

held-out accuracy; RF: legacy-refit representation-first flag. These secondary tests do not alter the primary verdict.

$D$	$p$	Fraction	$r$	Pass	$\bar{L}/\log p$	$\Delta\bar{L}$ (rose)	$\Delta A_{H}(W)$	Lead	$t_{90}$	Mis.
3	47	0.8	.45	10/10	0.52	$+9.5\%$ (10)	.51	2.4	19.5	0
3	47	0.7	.30	10/10	0.35	$+49\%$ (10)	.54	2.8	24.2	0
3	47	0.6	.18	10/10	0.15	$+233\%$ (10)	.48	3.7	34.1	0
3	47	0.5	.08	8/10	0.007	$+1117\%$ (10)	.48	4.7	48.1	0
3	61	0.8	.45	10/10	0.58	$+0.7\%$ (8)	.53	1.2	17.6	18
3	61	0.7	.30	10/10	0.43	$+13\%$ (10)	.63	1.3	22.0	1

Appendix

Table 3: Depths 3 and 4 at

p=47,61

across training fractions (10 fresh runs per setting;

\lambda_{V}=0.2\cdot 31/p

). Pass: the criterion of Appendix C.4 .

\bar{L}/\log p

: training CE (cycle mean) when held-out accuracy crosses

0.5

\Delta\bar{L}

: its relative change from the refit’s

0.2

crossing to the held-out

0.9

crossing, with the number of runs in which it rose.

\Delta A_{H}(W)

: last hidden layer. Lead: held-out minus refit

0.5

crossing.

t_{90}

: held-out

0.9

crossing. Mis.: the most training pairs misclassified at any log inside

W

, over all 10 runs. Medians over the 10 runs unless stated; updates in thousands.

Figure 10: Depth 3 at

p=61

across training fractions: every run (thin) and the median (thick) , in the format of Figure 4 (depth 3 at

p=47

); 10 fresh runs per fraction. Rows as in Figure 3 ; gray: from the median start to the median end of the post-hoc windows

W

. Table 3 gives

r

, the passes and the misclassified training pairs. Each column ends at its shortest run, so every run contributes at every plotted point; in this setting and in those of Figures 11 and 12 , the shortest run ends after the last run’s

90\%

held-out crossing.

Figure 11: Depth 4 at

p=47

across training fractions: every run (thin) and the median (thick) , as in Figure 10 .

Figure 12: Depth 4 at

p=61

across training fractions: every run (thin) and the median (thick) , as in Figure 10 .

Figure 13: Updating a hidden bias in different ways (one hidden layer,

p=31

, 10 initializations and splits; gray lines join the same run, bars are medians). Left: the relative decrease of

\bar{L}

from its largest post-knee value to the first update after which held-out accuracy stays at least

0.9

. Right: the captured share

q(I_{C})

of the primary endpoint, which requires at least

0.75

(dashed). Columns: no bias; the bias with the joint

[\boldsymbol{W}\mid b]

step; the joint step until the knee, then the bias frozen or only decaying; the bias updated by its own step throughout.

Setting	Runs	Pass	Fit	$\Delta\tilde{L}$	$q$	Final	Failing condition
Polar, gain $2\to 1$	10	10	200	$+33\%$ ( $+18$ , $+55$ )	0.93–1.00	0.999–1.000	–
Polar, no gain	10	0	1,190	$-8\%$ ( $-10$ , $-7$ )	0.35–0.51	0.999–1.000	share (10), refit (9)
Polar, fixed $g=1.5$	3	0	240	$-17\%$ ( $-18$ , $-16$ )	0.20–0.26	0.999–1.000	share, refit (3)
Polar, fixed $g=2$	3	0	200	$-27\%$ ( $-29$ , $-24$ )	0.13–0.17	1.000	share, refit (3)
Polar, fixed $g=3$	3	0	160	$-44\%$ ( $-49$ , $-42$ )	0.07–0.10	1.000	share, refit (3)

Appendix

Table 4: One block at

p=113

\phi=0.2

: the gain (seeds

i=11

–

20

; the 3-run rows use

i=11

–

13

). Fit: first full fit (median; updates).

\Delta\tilde{L}

: relative change of the averaged training CE from the no-mirror held-out

0.1

crossing to the

0.9

crossing (median and range).

q

: share of the all-pairs rise inside the all-pairs window. Final: held-out accuracy on all pairs.

	$\tau=2{,}500$	$5{,}000$	$10{,}000$	$20{,}000$	fixed
$g_{0}=1.5$	1/3, $+10\%$	3/3, $+13\%$	2/3, $+6\%$	2/3, $-4\%$	0/3, $-17\%$
$g_{0}=2$	3/3, $+28\%$	3/3, $+38\%$	3/3, $+24\%$	2/3, $-1\%$	0/3, $-27\%$
$g_{0}=3$	3/3, $+84\%$	3/3, $+104\%$	3/3, $+58\%$	2/3, $+9\%$	0/3, $-44\%$

Appendix

Table 5: The gain at

p=113

\phi=0.2

(seeds

i=11

–

13

): passes of 3 and the median change of the averaged training CE across the no-mirror held-out rise. The

g_{0}=2

\tau=5{,}000

entry uses the first three main runs. With

g_{0}=1.5

and

\tau=2{,}500

, all three runs pass the loose criterion.

Change at 12k	Pass	Window ends	First flip	Most	Final
None	0/10	30k (end)	17.8k–22.0k	45–438	0.941–0.994
Gain frozen ( $\eta_{g}=0$ , $g=1.30$ )	2/3	13.5k–13.7k	18.8k–19.7k	57–217	0.983–0.997
Head step $\eta_{V}$ : $0.03\to 0.01$	0/3	30k (end)	21.6k–23.0k	59–395	0.950–0.970
Feature step $\eta_{W}$ : $0.1\to 0.03$	2/3	12.0k	none	0	0.986–0.997

Appendix

Table 6: Two blocks at

\phi=0.2

(

\lambda_{V}=0.275

\tau=10{,}000

): one rate changed at update 12k (seeds

i=11

–

13

; no change:

i=11

–

20

). Window ends: end of the all-pairs window, in updates. First flip: first log after 12k with a misclassified training pair; Most: the most misclassified training pairs at a log after 12k. Final: held-out accuracy on all pairs.

$p$	$\phi$	$\lambda_{V}$	Training pairs	Pass	Loose	No-mirror	$\Delta\tilde{L}$	Final
61	0.30	0.45	1,116	1/5	1/5	1/5	$+87\%$	0.979–0.994
61	0.30	0.60	1,116	1/5	1/5	1/5	$+58\%$	0.964–0.991
61	0.33	0.45	1,227	5/10	5/10	10/10	$+106\%$	0.994–0.999
61	0.35	0.45	1,302	0/5	0/5	5/5	$+87\%$	0.998–0.999
61	$0.55^{\dagger}$	0.45	2,046	9/10	9/10	9/10	$+161\%$	0.9899–0.997
89	0.20	0.33	1,584	1/5	1/5	1/5	$+33\%$	0.39–0.996

Appendix

Table 7: One block across primes (gain

2\to 1

\tau=5{,}000

; seeds

i=11

–

15

, and

i=11

–

20

in the rows with 10 runs). Pass: strict, on both definitions; Loose: at most 10 misclassified training pairs; No-mirror: strict on the no-mirror definition alone, with the final all-pairs accuracy at least

0.99

\Delta\tilde{L}

: median change of the averaged training CE across the no-mirror rise, over the runs with a full rise. Final: held-out accuracy on all pairs. † Mirror-free split.

Readout	Crossing	$p=61^{\dagger}$	$p=61$	$p=89$	$p=113$	$p=149$
Refit, all frequencies	0.5	$-0.70$ (0)	$-0.68$ (0)	$-0.78$ (0)	$-1.31$ (0)	$-1.51$ (0)
	0.9	$-0.68$ (0)	$-0.72$ (0)	$-1.26$ (0)	$-2.15$ (0)	$-1.86$ (0)
Fourier refit, top 8	0.5	$+0.78$ (9)	$+0.80$ (10)	$+0.56$ (10)	$+0.86$ (10)	$+1.18$ (10)
	0.9	$-2.46$ (0)	$-1.16$ (0)	$+0.45$ (10)	$+0.90$ (10)	$+1.21$ (10)
Network’s logits, top 8	0.5	$+0.79$ (9)	$+0.76$ (10)	$+0.52$ (10)	$+0.81$ (10)	$+1.06$ (10)
	0.9	$-3.99$ (0)	$-1.70$ (0)	$+0.34$ (10)	$+0.74$ (10)	$+1.04$ (10)

Appendix

Table 8: Readout leads over the network on the no-mirror pairs: the network’s crossing minus the readout’s, in thousands of updates (positive: the readout crosses first), at the network’s

0.5

and

0.9

crossings; medians, with the number of leading runs in parentheses. Settings of Figure 14 , and

p=61

with the random split at

\phi=0.33

; † mirror-free split,

\phi=0.55

. Network’s logits, top 8: the network’s own centred logits projected, in class space, onto the 8 frequencies of the Fourier refit, without fitting.

Figure 14: One block across primes: every run (thin) and the median (thick) (gain

2\to 1

\tau=5{,}000

;

\lambda_{V}=0.45

0.33

0.275

and

0.23

; 10 runs per prime, seeds

i=11

–

20

). At

p=61

the training set is the mirror-free split of Appendix D.4 (fraction

0.55

), so no held-out pair has its mirror in it and the two held-out accuracies coincide. At each

p

the training fraction puts the start of the held-out rise after the CE minimum. Rows as in Figure 5 , plus the misclassified training pairs (symmetric log scale): the pass rule requires none inside the window, and the accuracy row cannot show a few pairs out of thousands. Gray: from the median start to the median end of the windows on all held-out pairs. Each column ends at its shortest run. Strict passes: 9/10, 10/10, 10/10 and 7/10; the run that fails at

p=61

ends at held-out accuracy

0.9899

, and at

p=149

all 10 runs pass the loose criterion. With a random split at fraction

0.33

p=61

passes in 5/10, and all five failures come from the mirror pairs (Figure 16 ).

Figure 15: Two blocks: every run (thin) and the median (thick) (training fraction

0.2

\lambda_{V}=0.275

, gain

2\to 1

with

\tau=10{,}000

; 10 runs, seeds

i=11

–

20

: the two-block runs of Figure 5 ). Panels: the rows of Figure 14 . No run passes: the CE rises during the held-out rise and no training pair is misclassified at any 20-update log inside the window until 17.8k–22.0k updates, when pairs begin to flip and held-out accuracy falls back. The three other two-block settings (5 runs each) and the runs of Table 6 are not shown.

Figure 16: Mirror pairs at

p=61

: the random split at

\phi=0.33

(left; 798 to 827 of the 2,494 held-out pairs have their mirror in the training set) and a mirror-free split at

\phi=0.55

(right; none), 10 runs each from the same initializations. Rows: training CE; held-out accuracy on all pairs (solid), the no-mirror pairs (dashed) and the mirror pairs (dotted), the refit and the Fourier-filtered refit on all held-out pairs, and training accuracy, medians (dot: median all-pairs accuracy when the window opens); where each accuracy rises (bar from the median

0.1

to the median

0.9

crossing, dot at the median

0.5

crossing, ticks at each run’s

0.5

crossing);

A_{H}

\hat{r}

. Gray: the all-pairs window (median). Bottom: the all-pairs share of the held-out rise inside the window against

a_{0}

, the all-pairs accuracy when the window opens, one point per run. Every run lies on the bound

(0.9-\max(a_{0},0.1))/0.8

; open symbols fail (the mirror-free run that fails, on its final accuracy, is hidden among the others at

a_{0}=0.02

Figure 17: Readouts against the network , on the no-mirror pairs. Top:

p=61

with the random split (

\phi=0.33

) and the mirror-free split (

\phi=0.55

): held-out accuracy of the network, of the refit (all frequencies) and of the Fourier-filtered refit with the top 8 and the top 16 frequencies, medians over 10 runs; below, where each rises (bars and dots as in Figure 16 ). Bottom: each readout’s lead over the network at the network’s

0.5

and

0.9

crossings (median over runs; bar from the smallest to the largest run) at the settings of Figure 14 and at

p=61

with the random split; above zero the readout crosses first. The top-16 readout is post hoc.

Neural Network Training Dynamics

Neural network training dynamics studies how parameters and representations change during optimization and how those changes shape learning, stability, and generalization.

See all 507 papers

A Systematic Empirical Study of Grokking: Depth, Architecture, Activation, and Regularization
Mar 26, 2026Shalima Binta Manir, Anamika Paul RupaGrokking Neural Network Training Dynamics
Howard University
Latent Algorithmic Structure Precedes Grokking: A Mechanistic Study of ReLU MLPs on Modular Arithmetic
Mar 24, 2026Anand SwaroopGrokking Neural Network Interpretability
A Spectral Theory of Grokking: Weight Decay induces Feature Learning
Sep 22, 2026Lenz Pracher, Pascal de Jong, Oskar Lieshaus +2Neural Network Generalization Weight Decay
Arnold-Sommerfeld-Center for Theoretical Physics, Ludwig-Maximilians-Universität München, Munich, Germany · Department of Applied Physics, Standford University, Standford, CA, USA · Department of Mathematics, University of Cambridge, Cambridge, UK

Mar 26, 2026J/K move · Enter open · S save

A Systematic Empirical Study of Grokking: Depth, Architecture, Activation, and Regularization

Shalima Binta Manir, Anamika Paul Rupa

Howard University

Grokking Neural Network Training Dynamics Neural Network Generalization Neural Network Optimization

Open paper AskPDF arXiv

Grokking the delayed transition from memorization to generalization in neural networks remains poorly understood, in part because prior empirical studies confound the roles of architecture, optimization, and regularization. We present a controlled study that systematically disentangles these factors on modular addition (mod 97), with matched and carefully tuned training regimes across models. Our central finding is that grokking dynamics are not primarily determined by architecture, but by interactions between optimization stability and regularization. Specifically, we show: (1) depth has a non-monotonic effect, with depth-4 MLPs consistently failing to grok while depth-8 residual networks recover generalization, demonstrating that depth requires architectural stabilization; (2) the apparent gap between Transformers and MLPs largely disappears (1.11

\times

delay) under matched hyperparameters, indicating that previously reported differences are largely due to optimizer and regularization confounds; (3) activation function effects are regime-dependent, with GELU up to 4.3

\times

faster than ReLU only when regularization permits memorization; and (4) weight decay is the dominant control parameter, exhibiting a narrow ``Goldilocks'' regime in which grokking occurs, while too little or too much prevents generalization. Across 3--5 seeds per configuration, these results provide a unified empirical account of grokking as an interaction-driven phenomenon. Our findings challenge architecture-centric interpretations and clarify how optimization and regularization jointly govern delayed generalization.

Abstract

Neural Network Training Dynamics

Explore similar work

A Systematic Empirical Study of Grokking: Depth, Architecture, Activation, and Regularization

Latent Algorithmic Structure Precedes Grokking: A Mechanistic Study of ReLU MLPs on Modular Arithmetic

A Spectral Theory of Grokking: Weight Decay induces Feature Learning