Subspace Crystallization in Grokking: Rank Compression Tracks the Phase Transition

Shlok Mehendale; Paras Chopra

cs.LGOct 26, 2026

Subspace Crystallization in Grokking: Rank Compression Tracks the Phase Transition

Authors: Shlok Mehendale, Paras Chopra

Organizations: Lossfunk

Abstract

Grokking describes understanding something deeply and intuitively. In neural networks, it is the phenomenon where a network generalizes long after memorizing its training data. This delay has been linked to a fall in the effective rank of the network's weight matrices, but whether this rank compression drives grokking or follows it is still unresolved. We study a mean-field MLP and a one-layer transformer trained on modular arithmetic and observe that: (i) Compression follows the onset, not memorization. Every transformer memorizes its training data within 250 steps, while the grokking onset (the first step at which test accuracy reaches 95%) varies from 6,500 to 21,000 steps. The steepest fall in effective rank comes about 900 steps after the onset, however late the onset occurs, and networks that memorize without grokking compress far less. (ii) Subspace crystallization. We call the span of the fewest top singular directions of a weight matrix that preserve 99% of test accuracy its sufficient subspace. Its size can shrink as the margin between the correct and the best wrong output grows, so we also count the directions needed to keep the outputs within 10%. In the transformer, this count falls during the transition from a quarter or more of the available directions to about a tenth, which we call subspace crystallization. (iii) Imposing low rank does not speed up grokking. Against a factorized full-rank control, capping the rank of the transformer's MLP input weights delays the onset, and capping the rank of its embedding stops 15 of 20 runs from grokking. Compression is therefore a signature of grokking rather than its cause.

Figures & tables

Panel of Figure 1 — Figure 1: Compression follows the grokking onset, and the directions needed to keep the outputs fall during the transition (one-layer transformer, 29 grokked seeds). (a) Every seed fits its training data by step 250, yet compression follows the grokking onset across a threefold range of onset times. (b) The number of directions needed to keep the outputs within 10% falls during the rise in test accuracy and settles at 8–13 per matrix.

	MLP	Transformer
1. Compression follows grokking, not memorization
Compression starts after test accuracy rises	45 of 45 runs	28 of 29 seeds
Fall in effective rank	10% if grokked	$60\pm 5\%$ if grokked, 15–23% if not
2. Subspace crystallization: the transformer’s solution contracts into a low-dimensional subspace as it groks
Directions needed per matrix	Accuracy: from 110 to 65 of 128 Outputs: 122–127 of 128 (MSE)	Outputs: from 27–53 to 8–13, during the transition
3. Imposing low rank does not speed up grokking

Table 1: Evidence for the three contributions in both architectures. Effective rank is measured on

W_{1}

(MLP) and

W_{\text{in}}

(transformer). Output counts are the fewest directions per matrix that change the network’s outputs by at most 10%.

Figure 2: The four training stages of the MLP (seed 0): rate of change of the effective rank of $W_{1}$ (top) and of test accuracy (bottom), computed from the five-step moving average used for the steepest-decrease measure (Appendix E.1 ). Stage 2 starts at the first step with test accuracy above 5%, Stage 3 at $t_{g}$ and Stage 4 at $t_{c}$ . Effective rank grows during memorization and the grokking transition, turns negative at $t_{g}$ and keeps falling through Stages 3 and 4.

Optimizer	Effective lr	$\\|W_{1}\\|_{F}$	$\sigma_{1}(W_{1})$	ER $(W_{1})$	$d_{0.99}$
Start (step 20,000)	–	726	111	115.1	65
AdamW, $\epsilon=10^{-6}$	$10^{3}$	658	85	117.8	65
Gradient descent	$10^{3}$	662	85	117.8	65
AdamW, $\epsilon=10^{-7}$	$10^{4}$	775	120	114.8	70
Gradient descent	$10^{4}$	783	126	114.6	72
AdamW, $\epsilon=10^{-8}$ (main)	$10^{5}$	1,040	203	104.8	89

Table 2: Stage 4 tracks the effective learning rate. Main seeds 0–2 are continued from step 20,000 to

10^{5}

with each optimizer, and entries are means. The effective learning rate is lr for gradient descent, lr

/(1-0.9)

for SGD with momentum 0.9, and lr

/\epsilon

for AdamW (

\text{lr}=10^{-3}

). For

\epsilon=10^{-9}

, lr

/\epsilon

is an upper bound, since

\sqrt{\hat{v}}

is comparable to

\epsilon

. Weight decay is set so that the per-step shrinkage matches AdamW’s.

d_{0.99}

: fewest directions that keep 99% of test accuracy.

Model	Checkpoint	Outputs within 10%	99% of accuracy	Median margin
MSE (15 seeds)	step 0 (untrained)	$116.4\pm 0.5$	–	$-0.004$
	step 11,000 (before onset)	$122.4\pm 0.5$	$119.3\pm 3.2$	$0.025$
	step 20,000	$122.3\pm 0.9$	$64.9\pm 1.7$	$0.470$
	end ( $10^{5}$ )	$127.3\pm 0.6$	$89.1\pm 2.3$	$0.611$
MSE, random labels (5)	end, training inputs	$123.0\pm 0.0$	–	$0.110$
CE, lr $3\times 10^{-3}$ (5)	last before onset	$121.4\pm 0.9$	$112.8\pm 5.5$	$3.8$

Table 3:

d_{\text{out}}

and

d_{0.99}

in the MLP (directions of

W_{1}

out of 128, on test inputs unless noted, mean

\pm

s.d.).

d_{\text{out}}

: fewest directions that keep the outputs within 10%.

d_{0.99}

: fewest directions that keep 99% of test accuracy. Cross-entropy (CE) models use centered logits.

Matrix	$r$	Grokked	$t_{g}$	Final test accuracy (%)
None (unfactorized)	–	5/5	$10{,}500\pm 1{,}777$	$100.0\pm 0.0$
$W_{\text{in}}$	128	5/5	$14{,}100\pm 4{,}095$	$99.6\pm 0.3$
	32	5/5	$14{,}350\pm 5{,}730$	$98.4\pm 1.6$
	24	4/5	$16{,}312\pm 3{,}332$	$95.5\pm 7.2$
	16	3/5	$16{,}417\pm 6{,}682$	$94.5\pm 7.7$
	8	1/5	$12{,}500$	$50.6\pm 39.7$

Table 4: Rank caps in the transformer (

W=AB

with inner dimension

r

, 5 seeds per arm,

t_{g}

averaged over the seeds that grok, mean

\pm

s.d.).

r=128

is the factorized control without a rank limit.

Step $-\,t_{g}$	$W_{E}$	$W_{\text{in}}$	$W_{\text{out}}$	$W_{U}$	Test acc. (%)	Restricted	Excluded
$-3{,}000$	$27.3\pm 2.8$	$28.0\pm 3.3$	$46.9\pm 7.6$	$52.6\pm 7.3$	$15.7$	$2.09$	$4.8$
$-2{,}000$	$26.0\pm 2.7$	$25.7\pm 3.3$	$40.2\pm 8.1$	$46.6\pm 7.9$	$26.2$	$1.37$	$8.1$
$-1{,}000$	$23.0\pm 3.1$	$21.0\pm 3.5$	$26.6\pm 7.4$	$33.4\pm 8.7$	$57.7$	$0.64$	$14.7$
$-500$	$19.9\pm 3.0$	$18.3\pm 2.9$	$16.7\pm 5.5$	$24.2\pm 8.0$	$82.4$	$0.26$	$16.6$
$0$	$17.2\pm 2.6$	$15.9\pm 2.3$	$11.0\pm 2.6$	$14.0\pm 4.9$	$96.9$	$0.12$	$16.2$
$+1{,}000$	$12.0\pm 2.2$	$13.9\pm 1.8$	$9.2\pm 1.9$	$9.1\pm 1.9$	$99.9$	$0.03$	$14.0$

Table 5: The transformer crystallizes during the transition (29 grokked seeds, each aligned at its own onset

t_{g}

, with checkpoints every 250 steps).

d_{\text{out}}

of each matrix (test inputs, mean

\pm

s.d.), test accuracy, and the restricted and excluded training losses of Nanda et al. (2023) (means).

Appendix figures & tables44 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Training dynamics of the MLP (seed 0): loss, weight norms, effective rank, accuracy, gradient norm, stable rank, L0 and Hoyer sparsity of $W_{1}$ , and Fourier IPR, against training step.

Variant	Grokked	$t_{g}$	Final ER	Stage-4 drop (%)	Fourier IPR
Main	5/5	$12{,}700\pm 447$	$104.9\pm 1.6$	$9.7\pm 1.5$	$0.457\pm 0.009$
wd 0	5/5	$12{,}700\pm 447$	$98.7\pm 2.2$	$14.9\pm 1.9$	$0.452\pm 0.010$
wd $10^{-3}$	5/5	$12{,}700\pm 447$	$99.3\pm 2.1$	$14.4\pm 1.9$	$0.453\pm 0.008$
wd $0.05$	5/5	$12{,}800\pm 274$	$113.7\pm 0.8$	$2.2\pm 0.4$	$0.466\pm 0.006$
wd $0.1$	5/5	$13{,}400\pm 418$	$116.7\pm 1.7$	$0.5\pm 0.4$	$0.476\pm 0.004$
CE, lr $10^{-3}$	1/5	$94{,}500$ ( $n=1$ )	$114.1\pm 0.3$	–	$0.200\pm 0.008$

Appendix

Table 6: MLP training variants (5 seeds each, mean

\pm

s.d.). Main: mean-squared error, quadratic activation, weight decay (wd)

10^{-2}

. CE: cross-entropy. Grokked: test accuracy

\geq 95\%

within

10^{5}

steps. Stage-4 drop: relative decrease of effective rank from its peak after

t_{g}

to its final value. Fourier IPR: Section 6.5 .

Panel of Figure 4 — Figure 4: Stage 4 across training variants (MLP, 5 seeds each, mean curves of test accuracy, effective rank of $W_{1}$ and, in (b, c), Fourier IPR).

$\lambda$	lr $\cdot\lambda$	Grokked	$t_{g}$	$d_{\text{out}}$ after $t_{g}$	$d_{\text{out}}$ , end	$d_{0.99}$ , end	ER $(W_{1})$ , end
0.01	$10^{-4}$	5/5	$10{,}100\pm 962$	$119.8\pm 0.4$	$69.6\pm 2.2$	$48.4\pm 2.5$	$71.1\pm 2.2$
0.03	$3\times 10^{-4}$	5/5	$4{,}500\pm 354$	$120.2\pm 0.8$	$100.4\pm 2.3$	$61.6\pm 3.2$	$98.7\pm 2.6$
0.1	$10^{-3}$	5/5	$2{,}300\pm 274$	$120.2\pm 0.8$	$118.4\pm 1.7$	$69.6\pm 2.6$	$115.7\pm 1.0$
0.3	$3\times 10^{-3}$	0/5	–	–	–	–	$1.0\pm 0.0$
1.0	$10^{-2}$	0/5	–	–	–	–	$1.0\pm 0.0$

Appendix

Table 7: Weight decay in the cross-entropy MLP (learning rate

10^{-2}

, 5 seeds each, mean

\pm

s.d.).

d_{\text{out}}

and

d_{0.99}

W_{1}

(of 128) on test inputs with centered logits, at the first checkpoint after

t_{g}

and at step

10^{5}

Step	20,000	30,000	50,000	100,000
$d_{0.99}$	$64.9\pm 1.7$	$69.3\pm 1.7$	$81.5\pm 2.4$	$89.1\pm 2.3$
ER $(W_{1})$	$114.4\pm 1.1$	$110.9\pm 1.5$	$107.0\pm 1.9$	$104.0\pm 2.1$
$\\|W_{1}\\|_{F}$	$731\pm 7$	$813\pm 14$	$937\pm 26$	$1059\pm 35$
$\\|W_{2}\\|_{F}$	$531\pm 5$	$597\pm 11$	$685\pm 21$	$763\pm 27$
$\sigma_{1}(W_{1})$	$111.9\pm 3.4$	$132.2\pm 3.0$	$171.1\pm 10.4$	$217.8\pm 21.6$
Size of top 65	$659\pm 10$	$748\pm 18$	$878\pm 29$	$1002\pm 37$

Appendix

Table 8: Why the sufficient subspace widens in Stage 4 (MLP, 15 seeds).

K=65

is the median

d_{0.99}

at step 20,000. Size:

(\sum_{i}\sigma_{i}^{2})^{1/2}

over the top 65 or the remaining 63 singular directions of

W_{1}

. Output change: median

\ell_{2}

change of the 97 outputs when the directions beyond the top 65 are removed, split into the two terms of

(h_{\text{top}}+h_{\text{rest}})^{2}-h_{\text{top}}^{2}

Weight decay	$d_{0.99}$	ER $(W_{1})$	$\\|W_{1}\\|_{F}$	$\sigma_{1}(W_{1})$	Size of remaining 63
0	$65.8\to 95.0$	$114.6\to 98.7$	$731\to 1{,}252$	$114\to 251$	$320\to 358$
$10^{-3}$	$65.4\to 93.8$	$114.7\to 99.3$	$731\to 1{,}220$	$113\to 249$	$320\to 354$
$10^{-2}$ (main)	$65.2\to 88.6$	$115.0\to 104.9$	$727\to 1{,}043$	$111\to 205$	$320\to 347$
$5\times 10^{-2}$	$64.0\to 73.0$	$115.7\to 113.7$	$710\to 825$	$101\to 151$	$321\to 342$
$10^{-1}$	$62.4\to 63.0$	$116.7\to 116.7$	$678\to 748$	$92\to 108$	$311\to 346$

Appendix

Table 9: Weight decay controls the widening (5 seeds each, step 20,000

\to

step

10^{5}

, means). All runs reach 100% test accuracy.

Figure 5: Stage-4 widening. Left: size of the top 65 and the remaining 63 singular directions of $W_{1}$ over training (15 seeds). Right: $d_{0.99}$ at step 20,000 and at the end of training for each weight decay (5 seeds).

	Step 20,000	Step 50,000	Step $10^{5}$
$\\|g\\|$	$3.4\times 10^{-6}$	$3.2\times 10^{-7}$	$1.9\times 10^{-7}$
$\langle g,\theta\rangle$ (= Euler expression)	$-6.0\times 10^{-4}$	$-2.5\times 10^{-4}$	$-2.2\times 10^{-4}$
$\cos(g,\theta)$	$-0.19\pm 0.01$	$-0.67\pm 0.07$	$-0.89\pm 0.06$
Share of coordinates with $\sqrt{\hat{v}}<\epsilon$	$0.870\pm 0.017$	$0.998\pm 0.001$	$0.9998\pm 0.0002$
Median $\sqrt{\hat{v}}$	$2.1\times 10^{-9}$	$7.3\times 10^{-10}$	$6.1\times 10^{-10}$
Mean $\|\hat{m}\|/(\sqrt{\hat{v}}+\epsilon)$ (1 = sign descent)	$0.21\pm 0.01$	$0.08\pm 0.01$	$0.07\pm 0.00$

Appendix

Table 10: Gradient and Adam state during Stage 4 (MLP, 15 seeds, mean

\pm

s.d.). Radial push of the update and decay pull

\lambda\|\theta\|^{2}

are per unit learning rate.

Optimizer	Effective lr	$\\|W_{1}\\|_{F}$	$\\|W_{2}\\|_{F}$	$\sigma_{1}(W_{1})$	ER $(W_{1})$	$d_{0.99}$
At step 20,000	–	726	529	111	115.1	65
AdamW, $\epsilon=10^{-10}$	$10^{7}$	1,793	718	346	95.7	99
AdamW, $\epsilon=10^{-9}$	$10^{6}$	1,582	874	318	94.7	98
AdamW, $\epsilon=10^{-8}$ (main)	$10^{5}$	1,040	748	203	104.8	89
AdamW, $\epsilon=10^{-7}$	$10^{4}$	775	551	120	114.8	70
AdamW, $\epsilon=10^{-6}$	$10^{3}$	658	469	85	117.8	65

Appendix

Table 11: Continuations from step 20,000 to

10^{5}

(main seeds 0–2, means). The effective learning rate is lr for gradient descent, lr

/\epsilon

for AdamW and lr

/(1-0.9)

for SGD with momentum.

Panel of Figure 6 — Figure 6: The sufficient subspace of the MLP (15 seeds). (a) Seed 0 and mean $\pm$ s.d. (b) $d_{0.9}$ and $d_{0.99}$ narrow sharply after the onset and widen again during Stage 4, while effective rank falls. $d_{\text{out}}$ does not narrow (Table 3 ).

Threshold $\tau$ (fraction of clean accuracy)	$d_{\tau}$ (mean $\pm$ s.d.)
90%	$80.5\pm 3.7$
95%	$84.5\pm 3.2$
99%	$89.1\pm 2.3$
99.9%	$93.1\pm 2.7$

Appendix

Table 12:

d_{\tau}

W_{1}

at four thresholds (step

10^{5}

, 15 seeds).

Step	Test acc. (%)	$d_{0.9}$	$d_{0.99}$	ER $(W_{1})$	Freq. for 90% power
10,000	$37.0\pm 8.0$	$97.5\pm 3.2$	$120.5\pm 3.5$	$115.7\pm 0.4$	$40.8\pm 0.4$
11,000	$67.3\pm 8.9$	$95.2\pm 2.1$	$119.3\pm 3.2$	$115.8\pm 0.4$	$40.3\pm 0.7$
12,500	$95.7\pm 2.7$	$81.7\pm 3.2$	$109.5\pm 3.2$	$115.9\pm 0.6$	$39.7\pm 0.5$
15,000	$100.0\pm 0.0$	$59.6\pm 1.8$	$77.0\pm 3.6$	$115.6\pm 0.8$	$38.5\pm 0.6$
20,000	$100.0\pm 0.0$	$50.8\pm 2.5$	$64.9\pm 1.7$	$114.4\pm 1.1$	$38.2\pm 1.1$
30,000	$100.0\pm 0.0$	$56.4\pm 2.4$	$69.3\pm 1.7$	$110.9\pm 1.5$	$38.3\pm 0.9$

Appendix

Table 13:

d_{0.9}

d_{0.99}

, effective rank and frequency usage of

W_{1}

over training (15 seeds).

Step $s$	Hidden space ( $\mathcal{C}$ )	Input space
0	$0.775\pm 0.003$	$0.608\pm 0.003$
5,000	$0.782\pm 0.003$	$0.623\pm 0.003$
10,000	$0.807\pm 0.005$	$0.677\pm 0.007$
12,500	$0.837\pm 0.007$	$0.723\pm 0.010$
15,000	$0.875\pm 0.010$	$0.773\pm 0.012$
20,000	$0.925\pm 0.013$	$0.843\pm 0.013$

Appendix

Table 14: Overlap of the top-89 singular subspaces of

W_{1}

at step

s

with those at step

10^{5}

(15 seeds). Hidden space: left singular vectors (the sufficient subspace

\mathcal{C}

). Input space: right singular vectors.

Matrix	$k=d_{0.99}$	Top- $k$	Bottom- $k$	Random subset	Random subspace
MLP $W_{1}$	$89.1\pm 2.3$ / 128	$99.26\pm 0.22$	$28.26\pm 6.18^{*}$	$31.26\pm 4.37$	$22.74\pm 3.93$
Transformer $W_{\text{in}}$	$13.0\pm 1.3$ / 128	$99.80\pm 0.28$	$5.50\pm 1.46^{\ddagger}$	$3.50\pm 0.54$	$12.49\pm 2.99$
Transformer $W_{\text{out}}$	$8.7\pm 1.5$ / 128	$99.95\pm 0.07$	$5.50\pm 1.46^{\ddagger}$	$4.42\pm 0.95$	$8.00\pm 1.93$
Transformer $W_{E}$	$7.7\pm 1.5$ / 97	$99.74\pm 0.35$	$1.01\pm 0.06$	$1.33\pm 0.23$	$1.31\pm 0.13$
Transformer $W_{U}$	$8.8\pm 1.4$ / 98	$99.86\pm 0.27$	$1.07\pm 0.32$	$1.80\pm 0.65$	$6.02\pm 1.92$

Appendix

Table 15: Top-

k

, bottom-

k

and random-

k

truncations at equal rank

k=d_{0.99}

(per seed). Test accuracy in %, mean

\pm

s.d. over 15 MLP seeds at step

10^{5}

and 14 grokked transformer seeds. Random controls average 10 draws. Chance is 1.03%.

Model	Checkpoint	$d_{0.99}$	$d_{\text{out}}$	Out05	Cen10	Marg10	Median margin
MSE, main	step 0	–	116.4	123.8	116.3	102.2	$-0.004$
	step 11,000	119.3	122.4	126.6	122.7	123.8	0.025
	step 20,000	64.9	122.3	125.7	122.4	113.1	0.470
	end	89.1	127.3	127.9	127.3	121.4	0.611
	end, training inputs	88.0	127.3	127.9	127.5	121.1	0.649
MSE, random labels	end, training inputs	–	123.0	127.0	123.2	116.6	0.110

Appendix

Table 16: Output counts in the MLP (directions of

W_{1}

out of 128, mean over seeds, test inputs unless noted).

d_{\text{out}}

and Out05: outputs within 10% and 5% (centered for CE). Cen10: per-input centered outputs within 10%. Marg10: per-input margins within 10%.

Model $A$ $\to$ grokked model $B$	Map	Test accuracy (%)
Grokked, other seed	fitted orthogonal	$95.82\pm 2.38$
Grokked, other seed	none	$1.03\pm 0.30$
Grokked, other seed	random orthogonal	$0.94\pm 0.27$
Random initialization	fitted orthogonal	$0.61\pm 0.12$
Random-label network	fitted orthogonal	$0.68\pm 0.07$
Grokked on $ab\bmod 97$ (other task)	fitted orthogonal	$1.42\pm 0.09$

Appendix

Table 17: Cross-seed transplants and linear CKA (MLP, final checkpoints unless noted, mean

\pm

s.d.). Transplant: test accuracy of model

B

’s readout applied to model

A

’s post-activations after the map (chance 1.03%). CKA on all

p^{2}

inputs.

Figure 7: Fourier power spectra of $W_{1}$ for six seeds (share of power per frequency, $f$ and $p-f$ combined).

Panel of Figure 8 — Figure 8: Phase structure and robustness of the MLP solution (seed 0).

	MLP (15 seeds)	Transformer (29 grokked)
Memorization step $t_{\text{fit}}$	$9{,}133\pm 297$	250 in all seeds
Grokking onset $t_{g}$	$12{,}633\pm 297$	$11{,}914\pm 4{,}116$
Start of compression $-\,t_{g}$	$t_{c}$ : $6{,}367\pm 3{,}248$ (15/15 $>0$ )	$t_{\text{on}}$ : $-1{,}888\pm 992$ (28/29 after $t_{10}$ )
Steepest ER decrease $-\,t_{g}$	$9{,}667\pm 3{,}320$ (15/15 $>0$ )	$905\pm 506$ (27/29 $>0$ )
Share of ER decrease after $t_{g}$	$0.99\pm 0.02$	$0.63\pm 0.09^{a}$

Appendix

Table 18: Timing of compression (mean

\pm

s.d., in steps).

t_{c}

is the first step at which ER falls below 99% of its running maximum.

t_{\text{on}}

and

t_{s}

are the start and the steepest point of the effective-rank drop of

W_{\text{in}}

after the initialization decay.

t_{10}

is the first step with test accuracy

\geq 10\%

Figure 9: Effective rank of $W_{\text{in}}$ (relative to step 0) and test accuracy of the 14 grokked transformer seeds against steps relative to $t_{g}$ . The effective rank falls steeply while the random initialization decays, plateaus at about 60% of its initial value, and drops again around $t_{g}$ .

Figure 10: Synthetic check of the steepest-drop measure (logging every 250 steps). Left: a sharp step at step 12,000 is located two logging intervals early, because the derivative of the centered average is tied over several steps and the earliest is taken. Right: a 1,000-step ramp centered on step 12,000 is located exactly.

Panel of Figure 11 — Figure 11: Compression lags grokking in every MLP seed.

Series	Null hypothesis	$\bar{F}$	$\bar{p}$	Med. $p$	Rejected ( $p<0.05$ )
First differences	Compression $\not\to$ Grokking	$2.84$	$0.133$	$0.059$	6/15
	Grokking $\not\to$ Compression	$24.77$	$0.003$	$<0.001$	15/15
Stationary (ADF)	Compression $\not\to$ Grokking	$2.44$	$0.257$	$0.126$	4/15
	Grokking $\not\to$ Compression	$19.77$	$0.003$	$<0.001$	15/15

Appendix

Table 19: Granger tests (temporal predictive asymmetry) across 15 MLP seeds.

Run	Final test acc. (%)	Peak rate (per 1,000 steps)	Relative change, step 4,000 $\to$ end
Grokked (29 seeds)	$\geq 95$	$-10.4\pm 2.0$	$-0.60\pm 0.05$
wd 1.0, seed 12	8.2	$-4.0$	$-0.15$
wd 0.1, seeds 0–2	1.8–4.3	$-2.1$ to $-2.3$	$-0.22$ to $-0.23$
wd 0.3, seeds 0–2	9.1–13.2	$-4.3$ to $-4.6$	$-0.17$ to $-0.19$

Appendix

Table 20: Effective rank of

W_{\text{in}}

after the initialization decay (transformer). Peak rate: steepest smoothed decrease after step 4,000.

Step	Group	$W_{E}$	$W_{\text{in}}$	$W_{\text{out}}$	$W_{U}$
10,000	before own onset (17)	$47.8\pm 2.4$	$77.2\pm 8.0$	$64.1\pm 4.5$	$60.4\pm 4.4$
	after own onset (12)	$39.0\pm 4.6$	$42.8\pm 12.3$	$33.2\pm 12.3$	$34.8\pm 7.1$
	never groks (1)	52	84	71	67
12,500	before own onset (9)	$44.3\pm 2.8$	$72.2\pm 7.4$	$61.7\pm 3.6$	$60.2\pm 3.7$
	after own onset (20)	$36.0\pm 4.1$	$32.1\pm 13.3$	$24.2\pm 15.1$	$30.4\pm 9.5$
	never groks (1)	50	86	70	66

Appendix

Table 21: Low rank is acquired at the transition (transformer, weight decay 1.0). Singular directions carrying 99% of

\sum_{i}\sigma_{i}^{2}

at the same training step, for grokked seeds before and after their own onset and for the seed that never groks. Mean

\pm

s.d., with the number of seeds in parentheses.

Figure 12: MLPs that memorize without generalizing barely compress. Effective rank of $W_{1}$ (left) and training accuracy (right) for 15 grokked MLPs (blue) and 5 random-label MLPs (red).

	Quantity	Transformer, held-out	MLP, $ab\bmod 97$	MLP, $(a{+}b)\bmod 113$
H1	$t_{s}-t_{g}$ ( $W_{\text{in}}$ )	$850\pm 524$ (14/15) ∗	–	–
H2	$t_{c}-t_{g}$	–	$6{,}300\pm 1{,}461$ (15/15) ∗	$6{,}700\pm 4{,}519$ (15/15) ∗
H3	$t_{s}-t_{g}$ ( $W_{1}$ )	–	$9{,}800\pm 2{,}274$ (15/15) ∗	$11{,}867\pm 6{,}741$ (15/15) ∗
H4	$d_{0.99}(s_{g})-d_{0.99}(s_{+})$	–	$32.9\pm 13.7$ (15/15) ∗	$10.8\pm 16.8$ (12/15) ∗
H5	$d_{0.99}(s_{\text{end}})-d_{0.99}(s_{+})$	–	$25.1\pm 2.9$ (15/15) ∗	$16.9\pm 3.1$ (15/15) ∗
H6	$t_{g}$ for $AB$ at $r=128/64/32$ , and unfactorized	–	$2{,}000/2{,}667/4{,}667$ and $12{,}833$	$2{,}000/3{,}000/4{,}500$ and $15{,}333$

Appendix

Table 22: Replication tests (one-sided tests, mean

\pm

s.d. in steps or directions, with the number of seeds showing a positive difference in parentheses).

	$ab\bmod 97$	$(a{+}b)\bmod 113$
Grokked	15/15	15/15
$t_{g}$	$12{,}633\pm 297$	$15{,}567\pm 417$
$t_{c}$	$18{,}933\pm 1{,}498$	$22{,}267\pm 4{,}629$
H2: $t_{c}-t_{g}$	$6{,}300\pm 1{,}461$	$6{,}700\pm 4{,}519$
	15/15, $t=16.70$ , $p=6.1\times 10^{-11}$	15/15, $t=5.74$ , $p=2.6\times 10^{-5}$
H3: $t_{s}-t_{g}$	$9{,}800\pm 2{,}274$	$11{,}867\pm 6{,}741$

Appendix

Table 23: H2–H5 on the two further MLP tasks (15 seeds each, mean

\pm

s.d.). Below each test, the number of seeds with a positive difference, the

t

statistic and the one-sided

p

Seed	$t_{g}$	$t_{s}-t_{g}$	Seed	$t_{g}$	$t_{s}-t_{g}$	Seed	$t_{g}$	$t_{s}-t_{g}$
15	7,000	$+500$	20	17,500	$+500$	25	17,750	$+250$
16	12,500	$+1{,}000$	21	8,000	$0$	26	8,750	$+500$
17	7,500	$+750$	22	13,000	$+1{,}000$	27	15,750	$+750$
18	17,000	$+1{,}000$	23	16,500	$+1{,}750$	28	8,500	$+750$
19	7,500	$+750$	24	10,250	$+2{,}000$	29	8,250	$+1{,}250$

Appendix

Table 24: H1 per seed: held-out transformer seeds,

t_{g}

and

t_{s}-t_{g}

(steps).

Model	Grokked	$t_{g}$	Final test acc. (%)	Final ER $(W_{1})$
Unfactorized $W_{1}$ (main)	3/3	$12{,}833\pm 289$	$100.0\pm 0.0$	$104.9\pm 1.1$
$W_{1}=AB$ , $r=128$ (control)	3/3	$2{,}000\pm 0$	$99.9\pm 0.2$	$89.1\pm 0.9$
$W_{1}=AB$ , $r=96$	3/3	$2{,}167\pm 289$	$100.0\pm 0.0$	$63.0\pm 5.7$
$W_{1}=AB$ , $r=80$	3/3	$2{,}500\pm 0$	$100.0\pm 0.0$	$49.4\pm 1.2$
$W_{1}=AB$ , $r=64$	3/3	$2{,}667\pm 289$	$100.0\pm 0.0$	$46.0\pm 3.7$
$W_{1}=AB$ , $r=48$	3/3	$3{,}167\pm 289$	$100.0\pm 0.0$	$38.0\pm 0.6$

Appendix

Table 25: Forcing low rank from initialization (MLP). Rank-constrained models use 3 seeds and the penalty 5 seeds.

Panel of Figure 13 — Figure 13: Forcing low rank from initialization (mean test accuracy and effective rank of $W_{1}$ against training step).

Step	$W_{1}=AB$ , $r=128$	Unfactorized $W_{1}$
0	$92.7\pm 0.3$	$114.6\pm 0.1$
500	$79.7\pm 1.0$	$114.5\pm 0.1$
1,000	$97.4\pm 0.6$	$114.5\pm 0.1$
1,500	$105.5\pm 0.3$	$114.6\pm 0.2$
2,000	$105.9\pm 0.3$	$114.7\pm 0.2$
3,000	$102.4\pm 2.0$	$114.9\pm 0.2$

Appendix

Table 26: The factorization biases

W_{1}

toward low rank. Effective rank of

W_{1}=AB

with

r=128

and of the unfactorized

W_{1}

(seeds 0–2,

(a{+}b)\bmod 97

). The factorized model groks at step 2,000, the unfactorized one at

12{,}833\pm 289

Task	Model	Grokked	$t_{g}$	Final test acc. (%)	Final ER $(W_{1})$
$ab\bmod 97$	Unfactorized $W_{1}$	3/3	$12{,}833\pm 289$	$99.7\pm 0.4$	$102.3\pm 0.7$
	$W_{1}=AB$ , $r=128$	3/3	$2{,}000\pm 0$	$99.3\pm 0.7$	$85.6\pm 1.4$
	$W_{1}=AB$ , $r=64$	3/3	$2{,}667\pm 289$	$100.0\pm 0.0$	$47.9\pm 0.2$
	$W_{1}=AB$ , $r=32$	3/3	$4{,}667\pm 289$	$99.0\pm 0.5$	$30.1\pm 0.2$
$(a{+}b)\bmod 113$	Unfactorized $W_{1}$	3/3	$15{,}333\pm 289$	$100.0\pm 0.0$	$109.5\pm 1.8$
	$W_{1}=AB$ , $r=128$	3/3	$2{,}000\pm 0$	$100.0\pm 0.0$	$89.9\pm 2.1$

Appendix

Table 27: Factorization runs on the two further tasks (H6) (3 seeds per arm, unfactorized: main seeds 0–2).

Task	Model	$t_{g}$	ER minimum before $t_{g}$ (step)	ER at $t_{g}$	ER at step 8,000
$(a{+}b)\bmod 97$	unfactorized	$12{,}483\pm 340$	–	$116.2\pm 0.2$	–
	$r=128$	$1{,}667\pm 58$	$76.9\pm 0.7$ ( $333$ )	$106.3\pm 0.5$	$90.7\pm 3.1$
	$r=64$	$2{,}450\pm 100$	$49.3\pm 0.2$ ( $267$ )	$60.5\pm 0.1$	$52.5\pm 6.1$
	$r=32$	$4{,}133\pm 29$	$28.2\pm 0.2$ ( $300$ )	$31.5\pm 0.1$	$31.0\pm 0.2$
$ab\bmod 97$	unfactorized	$12{,}567\pm 189$	–	$115.6\pm 0.6$	–
	$r=128$	$1{,}750\pm 0$	$73.8\pm 0.6$ ( $300$ )	$106.0\pm 1.3$	$87.3\pm 1.5$

Appendix

Table 28: Factorized MLPs logged every 50 steps (3 seeds each, mean

\pm

s.d.). ER: effective rank of

W_{1}=AB

Figure 14: Effective rank of factorized MLPs ( $W_{1}=AB$ with $r=128$ , 64 and 32, logged every 50 steps, 3 seeds each).

Model	Metric	$W$	Early	At window end	False alarms	Lead
MLP	Fourier IPR ( $W_{1}$ )	2,500	15/15	14/15	5/5	10,000
		5,000	15/15	15/15	5/5	7,500
		10,000	15/15	14/15	5/5	2,500
MLP	Gradient norm	2,500	15/15	0/15	5/5	9,000
		5,000	15/15	0/15	0/5	2,500
		10,000	3/15	0/15	3/5	0

Appendix

Table 29: Label-free onset detection (threshold

3\sigma

above the mean of the baseline window

[0,W)

). At window end : fires at the first logged step after the window. False alarms : fires in runs that never generalize (5 random-label MLP runs and 7 non-grokked transformer runs). Lead

=t_{g}-t_{\text{onset}}

(median, steps).

Figure 15: Transformer singular value spectra (seed 0) for $W_{E}$ , $W_{Q}$ , $W_{K}$ , $W_{V}$ , $W_{O}$ , $W_{\text{in}}$ , $W_{\text{out}}$ , $W_{U}$ and the MLP post-activations, with the 95% and 99% variance thresholds.

Figure 16: Fourier IPR of $W_{1}$ (blue) and test accuracy (red), MLP seed 0.

Matrix	Max rank	$k_{95}$	$k_{99}$	$d_{0.99}$
$W_{E}$ (embedding)	97	$8.8\pm 1.5$	$30.4\pm 4.0$	$7.7\pm 1.5$
$W_{Q}$ / $W_{K}$	32	$4.0\pm 0.0$	$4.0\pm 0.0$	$2.6\pm 0.5$ / $2.4\pm 0.7$
$W_{V}$	32	$18.5\pm 2.6$	$24.3\pm 2.6$	$10.9\pm 1.5$
$W_{O}$	128	$13.2\pm 1.4$	$21.3\pm 4.7$	$12.8\pm 1.3$
$W_{\text{in}}$	128	$13.1\pm 1.4$	$18.1\pm 3.8$	$13.0\pm 1.3$
$W_{\text{out}}$	128	$8.6\pm 1.4$	$11.8\pm 2.5$	$8.7\pm 1.5$

Appendix

Table 30: Transformer dimensionality (14 grokked seeds, final checkpoint).

k_{95}

k_{99}

: singular directions carrying 95% and 99% of

\sum_{i}\sigma_{i}^{2}

d_{0.99}

: fewest directions that keep 99% of test accuracy. Attention projections are stacked across the four heads. Activations are truncated after mean-centering.

Property	MLP (MSE)	MLP (CE)	Transformer
Loss, weight decay, training steps	MSE, $10^{-2}$ , $10^{5}$	CE, $10^{-2}$ , $10^{5}$	CE, $1.0$ , $2.5\times 10^{4}$
$d_{0.99}$ of first-layer weights (of 128)	$89.1\pm 2.3$	$48$ – $51$	$13.0\pm 1.3$
$d_{\text{out}}$ of first-layer weights (of 128)	$127.3\pm 0.6$	$70$ – $81$	$13.0\pm 1.9$
When $d_{\text{out}}$ falls	never	late, after $t_{g}$	during the transition
Frequencies holding 90% of first-layer power	$33.7\pm 1.9$ ( $W_{1}$ )	–	$4.2\pm 0.7$ ( $W_{E}$ )

Appendix

Table 31: Architecture comparison (final checkpoints, with 15 MLP seeds, 5 cross-entropy MLP seeds per learning rate and 14 grokked transformer seeds).

Run	Step	$W_{E}$	$W_{\text{in}}$	$W_{\text{out}}$	$W_{U}$	Test acc. (%)
wd 1.0, seed 12	5,000	31	33	64	70	2.8
	end	23	26	49	55	8.2
wd 0.1, seeds 0–2	5,000	45–47	42–43	74–77	81–84	1.0–2.8
	end	35–38	35–37	67–69	78–83	1.8–4.3
wd 0.3, seeds 0–2	5,000	40–42	38–40	71–73	79–82	1.2–2.9
	end	29–30	30–31	48–52	53–59	9.1–13.2

Appendix

Table 32: Runs that never generalize :

d_{\text{out}}

(centered logits within 10%, test inputs) at step 5,000 and at the end of training, with grokked seeds 3,000 steps after their onset for comparison.

Run	Train	Test	$W_{E}$	$W_{Q}$	$W_{K}$	$W_{V}$	$W_{O}$	$W_{\text{in}}$	$W_{\text{out}}$	$W_{U}$
wd 1.0, seed 12	100.0	8.2	42	4	4	30	42	68	62	62
wd 0.1, seed 0	100.0	4.3	43	4	19	31	50	82	78	72
wd 0.1, seed 1	100.0	1.8	43	4	18	31	51	81	78	73
wd 0.1, seed 2	100.0	3.0	44	4	18	31	52	84	79	73
wd 0.3, seed 0	100.0	13.2	34	4	4	31	43	74	61	62
wd 0.3, seed 1	100.0	10.7	36	4	4	31	43	73	61	62

Appendix

Table 33: Non-grokked transformer runs: final accuracy (%) and

k_{99}

(singular directions carrying 99% of

\sum_{i}\sigma_{i}^{2}

). Grokked seeds (mean):

W_{E}

30.4,

W_{Q}

4.0,

W_{K}

4.0,

W_{V}

24.3,

W_{O}

21.3,

W_{\text{in}}

18.1,

W_{\text{out}}

11.8,

W_{U}

23.8.

	Grokked	Not grokked
Subspace pair	wd 1.0	wd 0.1	wd 0.3	wd 1.0
	( $n=14$ )	( $n=3$ )	( $n=3$ )	( $n=1$ )
Attention write ( $W_{O}$ ) vs. MLP read ( $W_{\text{in}}$ )	$0.923\pm 0.035$	$0.839\pm 0.021$	$0.918\pm 0.048$	$0.947$
MLP read ( $W_{\text{in}}$ ) vs. MLP write ( $W_{\text{out}}$ )	$0.330\pm 0.050$	$0.269\pm 0.024$	$0.271\pm 0.043$	$0.218$
MLP write ( $W_{\text{out}}$ ) vs. unembedding read ( $W_{U}$ )	$0.885\pm 0.078$	–	–	–
MLP read vs. key frequencies of $W_{E}$	$0.252\pm 0.018$	–	–	–

Appendix

Table 34: Residual-stream alignment (mean cosine of principal angles between 11-dimensional subspaces). Not grokked: three seeds each at weight decay 0.1 and 0.3, and the one non-grokked seed at weight decay 1.0.

Panel of Figure 17 — Figure 17: Representational geometry (seed 0).

Grokking

Grokking is delayed generalization in which neural networks memorize training data before abruptly learning a rule that works on unseen examples.

See all 81 papers

Jul 7, 2026

At-Grok Is Not Converged:A Measurement-Validity Audit for Grokking Representation Metrics

On modular arithmetic, a network's embedding keeps compressing for tens of thousands of steps after it has already generalized. Reading effective rank at the grokking transition overstates the converged value by 3-5x on an MLP, and by 1.3-1.5x on a transformer trained to convergence; on the MLP it also erases which cells compress at all. Compression lags the accuracy transition by an amount on the order of the time-to-grok, at least 10,000 steps, rather than coinciding with it. A one-variable ablation shows what sets the lag size: adding LayerNorm to an otherwise identical transformer moves the fraction of compression done by the grok step from 0.87 to 0.25, and a pre-registered control rules out scale invariance as the mechanism. We package this as an audit that separates onset from compression, flags censoring, excludes boundary cells that never fully generalize, and checks that the reference floor has plateaued, with an adversarial suite that caught a false-confidence bug in our own branch. A secondary, MLP-specific depth law linking norm budget to converged floor fails a generality test on a transformer and flips sign under free weight decay. Code and the toolkit are released.

Truong Xuan Khanh

H&K Research Studio / Clevix LLC, Hanoi, Vietnam

Transformer Modular Arithmetic Grokking

Ask PDF

Feb 20, 2026

Global Low-Rank, Local Full-Rank: The Holographic Encoding of Learned Algorithms

Grokking -- the abrupt transition from memorization to generalization after extended training -- has been linked to the emergence of low-dimensional structure in learning dynamics. Yet neural network parameters inhabit extremely high-dimensional spaces. How can a low-dimensional learning process produce solutions that resist low-dimensional compression? We investigate this question in multi-task modular arithmetic, training shared-trunk Transformers with separate heads for addition, multiplication, and a quadratic operation modulo 97. Across three model scales (315K--2.2M parameters) and five weight decay settings, we compare three reconstruction methods: per-matrix SVD, joint cross-matrix SVD, and trajectory PCA. Across all conditions, grokking trajectories are confined to a 2--6 dimensional global subspace, while individual weight matrices remain effectively full-rank. Reconstruction from 3--5 trajectory PCs recovers over 95% of final accuracy, whereas both per-matrix and joint SVD fail at sub-full rank. Even when static decompositions capture most spectral energy, they destroy task-relevant structure. These results show that learned algorithms are encoded through dynamically coordinated updates spanning all matrices, rather than localized low-rank components. We term this the holographic encoding principle: grokked solutions are globally low-rank in the space of learning directions but locally full-rank in parameter space, with implications for compression, interpretability, and understanding how neural networks encode computation.

Yongzhong Xu

Neural Network Training Dynamics Grokking Low-Rank Approximation

Ask PDF

Feb 18, 2026

Low-Dimensional and Transversely Curved Optimization Dynamics in Grokking

Grokking -- the delayed transition from memorization to generalization in small algorithmic tasks -- remains poorly understood. We present a geometric analysis of optimization dynamics in transformers trained on modular arithmetic. PCA of attention weight trajectories reveals that training evolves predominantly within a low-dimensional execution subspace, with a single principal component capturing 68-83% of trajectory variance. To probe loss-landscape geometry, we measure commutator defects -- the non-commutativity of successive gradient steps -- and project them onto this learned subspace. We find that curvature grows sharply in directions orthogonal to the execution subspace while the trajectory remains largely confined to it. Importantly, curvature growth consistently precedes generalization across learning rates and hyperparameter regimes, with the lead time obeying a power law in the grokking timescale. Causal intervention experiments show that motion along the learned subspace is necessary for grokking, while artificially increasing curvature is insufficient. Together, these results support a geometric account in which grokking reflects escape from a metastable regime characterized by low-dimensional confinement and transverse curvature accumulation. All findings replicate across this learning-rate range, a qualitatively different slow regime (lr=5e-5, wd=0.1, 3 layers), and three random seeds, though alignment dynamics differ quantitatively between regimes. Causal intervention experiments establish that orthogonal gradient flow is necessary but not sufficient for grokking: suppressing it prevents generalization with a monotonic dose-response across four operations, while artificially boosting curvature defects has no effect.

Yongzhong Xu

Grokking Neural Network Training Dynamics Gradient Descent Dynamics Loss Landscape Geometry

Ask PDF

Abstract

Grokking

Explore similar work

At-Grok Is Not Converged:A Measurement-Validity Audit for Grokking Representation Metrics

Global Low-Rank, Local Full-Rank: The Holographic Encoding of Learned Algorithms

Low-Dimensional and Transversely Curved Optimization Dynamics in Grokking