Is Grokking a Computational Glass Relaxation?

TL;DR

This work reinterprets the grokking phenomenon as a computational glass relaxation, framing neural networks as physical systems where training loss acts as energy and parameter configurations define entropy. By sampling the Boltzmann entropy landscape $S(\ln(L_{train}), A_{test})$ with Wang-Landau Molecular Dynamics, the authors show there is no entropy barrier between memorization and generalization, arguing against a first-order phase transition and instead a slow relaxation toward higher-entropy generalizing states. They demonstrate a pronounced high-entropy advantage in grokking tasks and show that constraining weight norms reduces but does not remove this advantage, linking generalization to entropy rather than weight decay alone. Finally, they introduce WanD, a physics-inspired optimizer based on WLMD that can achieve generalization with comparable efficiency to AdamW while largely eliminating grokking, highlighting the potential for entropy-guided optimizer design in improving generalization and avoiding non-generalizable glassy states.

Abstract

Understanding neural network's (NN) generalizability remains a central question in deep learning research. The special phenomenon of grokking, where NNs abruptly generalize long after the training performance reaches a near-perfect level, offers a unique window to investigate the underlying mechanisms of NNs' generalizability. Here we propose an interpretation for grokking by framing it as a computational glass relaxation: viewing NNs as a physical system where parameters are the degrees of freedom and train loss is the system energy, we find memorization process resembles a rapid cooling of liquid into non-equilibrium glassy state at low temperature and the later generalization is like a slow relaxation towards a more stable configuration. This mapping enables us to sample NNs' Boltzmann entropy (states of density) landscape as a function of training loss and test accuracy. Our experiments in transformers on arithmetic tasks suggests that there is NO entropy barrier in the memorization-to-generalization transition of grokking, challenging previous theory that defines grokking as a first-order phase transition. We identify a high-entropy advantage under grokking, an extension of prior work linking entropy to generalizability but much more significant. Inspired by grokking's far-from-equilibrium nature, we develop a toy optimizer WanD based on Wang-landau molecular dynamics, which can eliminate grokking without any constraints and find high-norm generalizing solutions. This provides strictly-defined counterexamples to theory attributing grokking solely to weight norm evolution towards the Goldilocks zone and also suggests new potential ways for optimizer design.

Figures (8)

• Figure 1: Entropy landscape of three different arithmetic tasks : (a) $x+y\ mod\ 67$ is easy to learn, max entropy state can generalize better than training line. (b) $x^2+y\ mod\ 67$ is hard to learn, the generalization advantage of the max entropy state over the training results under the same training loss is very significant. (c) $x^3+xy^2+y\ mod\ 67$ cannot be learned by training, max entropy state also can not generalize.
• Figure 2: Entropy landscape of modular addition ($p=67$) task, sampling in three transformer models that have different widths. Here weight norm is constrained to a fixed value of 30. Compared with the grokking NNs, the high-entropy advantage is significantly weakened but still exists.
• Figure 3: Using (a)(b)(c) WanD and (d) AdamW to train a one layer Transformer on modular addition ($p=67$) task. (a) Train and test accuracy over time. The grokking is largely eliminated while training with WanD. (b) Weight norm over time. (c) Test accuracy over weight norm. These (a, b, c) together demonstrate that the model can generalize at high weight norm. (d) Train and test accuracy over time. Both WanD and AdamW use the same learning rate.
• Figure 4: Visualization of three different arithmetic tasks. (a) $x+y\ mod\ 67$. (b) $x^2+y\ mod\ 67$. (c) $x^3+xy^2+y\ mod\ 67$. Our protocol is equivalent to randomly taking 50% of the pixels from the image as the training set and trying to complete the rest. The more complex the expression, the more visually irregular the image is.
• Figure 5: Training results of three different arithmetic tasks: (a) $x+y\ mod\ 67$ is easy to learn, but the hysteresis of the test accuracy shows grokking. (b) $x^2+y\ mod\ 67$ is hard to learn, the grokking is more obvious. (c) $x^3+xy^2+y\ mod\ 67$ cannot be learned.