A New Take on Grokking: Rank Compression Marks the Phase Transition, Not Its Cause

The "grokking" phenomenon in neural networks—where a model suddenly generalizes after memorizing training data for a long time—has long fascinated researchers. Some previously argued that the decline in the effective rank of weight matrices is a key factor driving grokking, but a new study, through systematic experiments, gives a different answer: compression follows grokking rather than causing it.

Key Finding: Compression Lags Behind the Onset of Grokking

The research team trained mean-field MLPs and single-layer Transformers on modular arithmetic tasks. They observed that all Transformers memorized the training set within 250 steps, but the onset of grokking (the step at which test accuracy first reached 95%) ranged from 6,500 to 21,000 steps. Interestingly, no matter how late grokking occurred, the steepest decline in effective rank always appeared about 900 steps after the onset of grokking. Networks that only memorized without generalizing showed much less compression.

These data directly challenge the hypothesis that "rank compression drives generalization." If compression were the cause, it should appear before grokking—but the facts are exactly the opposite.

Effective rank is a metric for the "actual dimensionality in use" of a matrix, typically defined from the singular value distribution: either as the exponential of the entropy after normalizing singular values into a probability distribution, or as the number of directions needed to cover a given proportion of energy. When the effective rank of a weight matrix decreases, it means the network concentrates information into a few singular directions rather than spreading it evenly across all dimensions. In neural networks, effective rank is often used as a proxy for representation complexity or compression, leading some researchers to suspect it drives generalization. But a decrease in effective rank is itself only a descriptive geometric change and does not directly establish causal direction; it must be judged against the training timeline to determine its temporal relationship with the generalization phase transition.

Subspace Crystallization: From a Quarter to a Tenth

The study also introduces a more refined concept: subspace crystallization. They define the "sufficient subspace" as the space spanned by the minimal set of top singular directions of the weight matrix that retains 99% of test accuracy. Because the size of this subspace shrinks as the margin between the correct output and the best incorrect output grows, the researchers also counted the number of directions needed to keep output error within 10%.

In Transformers, this number of directions dropped during the phase transition from more than a quarter of the available directions to about a tenth. The vivid analogy is that the network seems to "crystallize" a low-dimensional subspace during grokking, condensing key information into a few directions.

Singular value decomposition (SVD) decomposes any matrix into the product of two orthogonal matrices and a diagonal singular value matrix, with singular values ordered from largest to smallest; the corresponding left and right singular vectors span the principal directions of the matrix's row and column spaces. When analyzing weight matrices, the top singular directions correspond to the input and output directions to which the network is most sensitive and which carry the most information; discarding directions with small singular values amounts to a low-rank approximation. The phrase "the subspace spanned by the top singular directions" in the text is defined precisely via SVD: keeping the first k singular directions approximates the original matrix. The subspace shrinking as the output margin grows means that a generalized network can retain key functionality with fewer directions—this is the geometric meaning of "crystallization."

Forcing Low Rank Actually Slows Down Grokking

To test causality, the researchers set up control experiments: imposing a rank cap on the MLP input weights of the Transformer delayed the onset of grokking compared with a full-rank decomposition control group; imposing a rank cap on the embedding layer completely prevented grokking in 15 out of 20 runs.

In other words, artificially lowering the rank does not accelerate generalization and may instead hinder grokking. This further supports the conclusion that "compression is an accompanying feature of grokking."

Modular arithmetic tasks (such as modular addition and modular multiplication) are classic benchmarks for studying grokking: the input is two integers and an operation, and the output is the result modulo some prime. Such tasks have a discrete algebraic structure, and the network must first memorize the training samples before suddenly learning a generalizable algorithmic rule late in training. The sample space of modular arithmetic is finite, making it easy to precisely control the training/test split, and generalization behavior is sensitive to architecture and hyperparameters, so it is often used to observe phase-transition signals such as the onset of grokking and changes in effective rank. Comparing mean-field MLPs and single-layer Transformers on such tasks helps distinguish the effects of architectural factors from those of optimization dynamics on grokking.

Implications for Understanding Generalization

This study repositions grokking as a phase transition process, with rank compression and subspace crystallization as hallmark signals of this transition. For practitioners, this means monitoring effective rank can serve as an indicator of whether a model has entered the generalization stage, but one should not expect to "induce" grokking by forcing low rank. Understanding the underlying mechanism of generalization still requires deeper exploration from the perspectives of optimization dynamics and representation learning.

One-line Summary

A new study shows that rank compression in weight matrices occurs after grokking—it is an accompanying signal, not a driving factor.

Paragraph Summary

A new study, through systematic experiments on modular arithmetic tasks, re-examines the role of effective rank decline in the grokking phenomenon in neural networks. The study finds that Transformers memorize the training set within 250 steps, but the onset of grokking ranges from 6,500 to 21,000 steps, while the steepest decline in effective rank always appears about 900 steps after the onset of grokking; networks that only memorize without generalizing show less compression. The study further proposes the concept of "subspace crystallization": during the phase transition, the number of directions needed to retain 99% of test accuracy drops from more than a quarter of the available directions to about a tenth. Causal control experiments show that forcing low rank on the MLP input weights delays grokking, while forcing low rank on the embedding layer completely prevents grokking in most runs. The conclusion is that compression is an accompanying feature of grokking rather than a driving factor; monitoring effective rank can serve as an indicator of the generalization stage, but grokking cannot be induced by forcing low rank.