The Question

Grokking, memorization followed by delayed generalization, is a rare clean window into how circuits emerge from gradient descent. We know a lot about what grokking networks compute: a sparse Fourier algorithm for modular arithmetic, formed through phases of memorization, circuit formation and cleanup. A recurring observation is that the effective rank of their weight matrices falls as test accuracy jumps. We ask two questions.

The usual evidence cannot settle these, for three reasons:

1. Timing is ambiguous. Compression is usually seen after both memorization and the rise in test accuracy, so it is unclear which of the two it follows.
2. Timing alone cannot show cause. That needs interventions.
3. Accuracy hides the margin. Truncating a weight matrix until test accuracy drops only checks which output is largest. The margin is the gap between the correct output and the largest wrong one. As margins grow, more directions can be deleted before any prediction flips, so the count can fall even though the outputs still depend on those directions.

Setup

We train two architectures on $(a+b) \bmod 97$ across many random seeds, and replicate on $ab \bmod 97$ and $(a+b) \bmod 113$. The grokking onset $t_g$ is the first step with test accuracy $\geq 95\%$. The memorization step $t_{\text{fit}}$ is the first step with training accuracy $\geq 99\%$.

Counting Directions

Our primary geometric metric is the effective rank, the exponential of the entropy of the normalized singular values: $$ \mathrm{ER}(W) = \exp\big(H(\hat\sigma)\big), \qquad \hat\sigma_i = \sigma_i \Big/ \textstyle\sum_j \sigma_j . $$ To ask how many directions the solution needs, we replace a weight matrix $W = U\Sigma V^\top$ by its top-$k$ truncation $W_{[k]} = U_{:k}\Sigma_{:k}V_{:k}^\top$ and leave the rest of the network unchanged. Two counts follow:

• $d_{0.99}$: the fewest directions that keep 99% of the test accuracy. Their span is the sufficient subspace.
• $d_{\text{out}}$: the fewest directions that keep the network's outputs within 10%, i.e. $\lVert f_{[k]} - f \rVert_F \leq 0.1\,\lVert f \rVert_F$ over the test inputs.

$d_{0.99}$ can fall simply because margins grow. $d_{\text{out}}$ cannot. We call a narrowing of $d_{\text{out}}$ subspace crystallization. The 10% tolerance is a choice. At 5%, the MLP's counts tell the same story (Table 16 of the paper).

Four Stages of MLP Training

The rates of change of effective rank and test accuracy split the MLP's training into four stages. In memorization, training accuracy reaches 97.8% by step 8,000 while test accuracy stays below 5%. The grokking transition follows, with the onset at step 12,633 ± 297. In consolidation, effective rank stays near its maximum while the margin grows. In late compression, 99% of the total fall in effective rank happens. In the MLP, rank grows during memorization and starts to fall only once test accuracy has risen. Seeing the last stage takes training well past the point where test accuracy saturates.

Compression Tracks the Grokking Onset

The transformer pulls memorization and generalization apart. All 29 grokked seeds fit their training set within the first 250 steps, yet their onsets range from 6,500 to 21,000 steps. Weight decay first shrinks the random initialization, which by itself lowers the effective rank of $W_{\text{in}}$ to a plateau, so we measure compression only after step 4,000. A second, sharp drop follows, and the step $t_s$ at which it is steepest tracks the onset almost exactly: Pearson $r = 0.99$, slope 1.02 (95% CI 0.98–1.07). It scatters by 506 steps around $t_g$ but by 4,234 steps around the memorization step.

Networks that memorize without grokking compress far less:

Both transformer falls are measured from step 4,000. Six of the seven runs that never grok used weaker weight decay (0.1 or 0.3), which alone means less shrinkage. A comparison at matched steps avoids this. Before its own onset, each grokked seed has as many directions as the one non-grokking run at weight decay 1.0. After its onset, it has far fewer.

Compression Does Not Start Before Test Accuracy Rises

In the MLP, compression starts after the onset in all 15 seeds, with a mean lag of 6,367 ± 3,248 steps (one-sided $t = 7.59$, $p = 1.3 \times 10^{-6}$). Granger tests agree: changes in test accuracy predict changes in effective rank in all 15 seeds, and the reverse holds in at most 6.

In the transformer, compression is part of the transition itself. The second drop begins 1,888 ± 992 steps before the onset, but only after test accuracy has already reached 10% in 28 of 29 seeds. It is steepest 905 ± 506 steps after the onset, and 63% of it happens after.

Replication on Held-Out Seeds and Tasks

We fixed the timing measures on seeds 0–14 and then tested them on data they had not seen. On 15 held-out transformer seeds, the steepest decrease came 850 ± 524 steps after the onset (14 of 15 seeds, the fifteenth at the onset), close to the 964 ± 499 steps on the original seeds. On $ab \bmod 97$ and $(a+b) \bmod 113$, the MLP's compression followed the onset in all 30 runs. With a Bonferroni correction over all nine $t$-tests, every test but one remains significant.

Subspace Crystallization

With a checkpoint every 250 steps and each seed aligned at its own onset, we can watch $d_{\text{out}}$ narrow while test accuracy rises. Each square below is one singular direction of a transformer weight matrix, and the filled squares are the leading directions needed to keep the network's outputs within 10%. Drag the slider or press play.

Three thousand steps before the onset, at 16% test accuracy, the outputs depend on 22–54% of each matrix's directions ($d_{\text{out}}$ = 27–53). At the onset this is 9–18%, and 3,000 steps later 7–11% ($d_{\text{out}}$ = 9–14). For $W_{\text{in}}$, $W_{\text{out}}$ and $W_U$ it is halfway down 740–960 steps before the onset, at about 60% test accuracy, in 28 or 29 of 29 seeds.

It Coincides With Cleanup

The narrowing lines up with the cleanup phase of Nanda et al. (2023). Their restricted loss is halfway down 1,966 ± 647 steps before the onset, as the generalizing circuit forms. Their excluded loss is halfway up 1,026 ± 646 steps before the onset, as the memorizing components are removed, at the same time as $d_{\text{out}}$ narrows. $d_{\text{out}}$ adds that cleanup reduces the directions the outputs depend on, not only those the predictions depend on. Transformers that never grok narrow only slowly and stay two to nine times wider.

The Final Solution Uses Few Directions

At the final checkpoint, $d_{0.99}$ is 7.9% of the directions of $W_E$ (7.7 of 97), 6.8% of $W_{\text{out}}$ (8.7 of 128) and 10.2% of $W_{\text{in}}$ (13.0 of 128), and the MLP's post-activations need only 3.4% (17 of 512). These counts are not inflated by the margin: $d_{\text{out}}$ is 8.1–13.0 directions for all four matrices, roughly 6–13% of each. The leading directions really carry the solution: truncating to them keeps at least 99.26% of test accuracy, while random subspaces of the same rank keep at most 12.5%.

Judging Truncation by Accuracy Can Show Compression That Isn't There

The MLP trained with mean-squared error shows why the counting criterion matters. Its sufficient subspace narrows right after the onset, from $d_{0.99} = 109.5$ at step 12,500 to $64.9$ at step 20,000, then widens again to $89.1$. But over the same stretch the median margin grows more than fourfold, and $d_{\text{out}}$ does not fall at all. The apparent narrowing measured how far predictions were from flipping, not how many directions the function uses.

With cross-entropy, the MLP does crystallize, but late, and as part of a different process: $d_{\text{out}} \approx 120$ just before and after the onset, ending at 70–81 and tracking the late fall in effective rank rather than the transition.

Why the Transformer Crystallizes Further

We attribute the gap to the architecture. The transformer's embedding uses about four frequencies (4.2 hold 90% of its Fourier power), and each frequency needs two directions, a cosine and a sine. That predicts about 8 directions, close to $d_{0.99}$ of $W_E$, $W_{\text{out}}$ and $W_U$ (7.7–8.8). The directions attention writes to nearly coincide with those $W_{\text{in}}$ reads from (mean principal-angle cosine 0.923, against 0.248 for random subspaces), so the MLP block receives an input already confined to a few directions. The MLP's first layer instead reads both one-hot operands directly, so each frequency costs it four directions, and it spreads its weights over most of the 48 frequencies.

Forcing Low Rank

If compression drove grokking, removing directions in advance should help. We write a weight matrix as $W = AB$ with inner dimension $r$ and train from initialization. The comparison that matters is with a factorized control at full rank ($r = 128$), because the factorization by itself changes when grokking happens.

Rank caps in the transformer, with 5 seeds per arm. $t_g$ is averaged over the seeds that grok.

The trained solution needs only about 13 directions of $W_{\text{in}}$ and 8 of $W_E$, so caps above these leave its capacity intact. Even so, capping $W_{\text{in}}$ at 32 leaves the onset where the control has it, capping it at 24 or 16 delays it and makes it less reliable, and capping it at 8 prevents grokking in four of five seeds. Capping $W_E$ is worse: only five of twenty capped runs grok.

Why the Control Matters

In the MLP, factorizing $W_1 = AB$ at full rank speeds grokking up seven- to eight-fold on its own ($t_g$ = 1,667 against 12,483 steps), whereas in the transformer it slows grokking down. Both fit known mechanisms: factorization preconditions the updates of the product, and weight decay on $A$ and $B$ penalizes the nuclear norm of $W$, since $\lVert A\rVert_F^2 + \lVert B\rVert_F^2 \geq 2\lVert AB\rVert_*$, a penalty that is negligible at the MLP's weight decay and strong at the transformer's. Against the factorized control, MLP rank limits delay grokking. Compared with the unfactorized network alone, low rank would wrongly appear to accelerate it.

A nuclear-norm penalty tells the same story. Applied to $W_1$ from step 0, it compresses the matrix well before grokking (effective rank 78.5 against 115.9 at step 11,500) and delays the onset from 12,700 to 16,200 steps.

More Findings

The MLP's Late Compression Is Set by AdamW's ε

Late in training, the square root of Adam's second-moment estimate falls below $\epsilon$ for 87–100% of coordinates, so the update $m/(\sqrt{v}+\epsilon)$ becomes $m/\epsilon$: AdamW acts as gradient descent with momentum at learning rate $\mathrm{lr}/\epsilon = 10^5$. Continuing training with other optimizers confirms it is this effective learning rate, not Adam itself, that governs the late stage. SGD with momentum at the same effective rate reproduces all of it. Concurrent work by Barron et al. (2026) reaches the same small-gradient regime from the cross-entropy side. There, the balance between Adam's $\epsilon$ and weight decay sets the margin, at roughly $\log(1/\epsilon)$.

Same Solution, Different Coordinates

Different MLP seeds favour different Fourier frequencies, and their top weight subspaces overlap no more than random ones. Yet after an orthogonal Procrustes rotation, one seed's readout decodes another seed's hidden activations with 95.8% accuracy, against 1.0% without the rotation. The rotation does not rescue a random network, a network trained on random labels, or one grokked on a different task (0.6–1.4%).

Fourier IPR Is Not an Early-Warning Signal

The Fourier inverse participation ratio, which measures how periodic the hidden units are, rises before the onset. But it also rises in every network that never groks: all five random-label MLPs and all seven non-grokking transformers. It tracks the progress of training rather than an impending transition.

Takeaways

Compression tracks the grokking onset, not memorization. Forcing low rank from initialization does not reliably bring grokking forward. Compression is a signature of grokking, not a driver. Three results matter beyond grokking, for anyone who measures how many directions a network uses:

1. Judging truncation by accuracy can show compression that isn't there. In the MSE-trained MLP, $d_{0.99}$ fell from 110 to 65 of 128 directions after the onset, while $d_{\text{out}}$ stayed at 122–127. The narrowing was the margin growing. $d_{0.99}$ later rose again to 89, so it is not even a fixed property of the trained model.
2. A rank cap is also a reparameterization, and the baseline flips the conclusion. Writing $W = AB$ at full rank made the MLP grok 7.5x sooner on its own. Against the unfactorized network, rank caps look 3–5x faster. Against the factorized control, they are 1.5–2.5x slower.
3. Spectral signals need negative controls. Fourier IPR rose in all 5 random-label MLPs and all 7 transformers that never grokked, so it cannot flag an upcoming onset. Effective rank falls in those runs too, by 15–23% against 60% in runs that grok, and it drops early in every run as weight decay shrinks the random initialization. Only the comparison with controls separates the grokking signal from these.

Limitations

All our tasks are modular arithmetic on cyclic groups. Non-abelian groups, and language or vision models, would test the results more broadly. The two architectures also differ in loss and training length. We varied weight decay separately, but the cross-entropy MLPs need larger learning rates to grok. Granger tests capture only temporal order, so our causal conclusions rest on the rank interventions.

Citation

If you find this work useful, please cite:

@misc{mehendale2026subspace,
  title         = {Subspace Crystallization in Grokking: Rank Compression
                   Tracks the Phase Transition},
  author        = {Mehendale, Shlok and Chopra, Paras},
  year          = {2026},
  eprint        = {2610.35112},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2610.35112}
}

Code to reproduce every experiment is at github.com/Lossfunk/Causal-geometric-grokking. Correspondence: shlok.mehendale@lossfunk.com.