Weight matrix compression follows deep learning insight in neural networks
Subspace Crystallization in Grokking: Rank Compression Tracks the Phase Transition
Machine LearningArtificial Intelligence
Summary
Grokking in machine learning means a model starts to understand a problem well long after simply memorizing the data. The authors studied neural networks and found that changes in the complexity of certain weight matrices happen after the model begins to generalize, not before. They identified a pattern called subspace crystallization, where fewer key directions in the weight matrix are needed to keep performance high. Finally, they showed that forcing the model to compress these matrices early does not speed up grokking, meaning the compression is a sign of learning, not the cause.
What this means in practice
- •For machine learning engineers: Design training schedules that expect rank compression as a sign of learning progress rather than a prerequisite for early generalization.
- •For ai system integrators: Avoid premature model simplification by constraining weight matrix ranks early, since it can delay or prevent deep understanding in neural networks.
Authors
Shlok Mehendale, Paras Chopra
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.