Is Grokking a Computational Glass Relaxation?
Authors: Xiaotian Zhang
Yue Shang
Entao Yang
Ge Zhang
Venue: NeurIPS 2025 spotlight ()
Keywords: grokking
entropy
generalization
glass
Primary area: theory
TL;DR: By framing grokking as computational glass relaxation, this work explains grokking from the perspective of Boltzmann entropy and proposes a physics-based grokking-resistant optimizer.

Abstract: Understanding neural network' (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 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.

Is Grokking a Computational Glass Relaxation?
Xiaotian Zhang'
Yue Shang!: 1.2
Entao Yang". Ge Zhang!.
Department of Physics, City University of Hong Kong, Hong Kong. China Department of Physics and Astronomy, University of Pennsylvania, Philadelphia, PA, USA Innovation Campus Delaware, Air Liquide, Newark, DE, USA
Abstract
Understanding neural network's (NN) generalizability remains 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 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 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 (density of states) landscape as a function training loss and test accuracy. Our experiments in transformers there is NO entropy barrier in the memorization- tooooottttoggeeee to-generalizztion transition grokking, challenging previous theory that defines grokking first-order phase transition M. We identify a high-entropy advantage under grokking, an extension of prior work linking entropy generalizability but much more significant [2]. Inspired by grokking's ' far-from- rrom-equilibrium nature, we develop a toy optimizer, WanD, based on Wang-Land Molecular Dynamics, which can eliminate grokking without any constraints and find high-norm generalizing solutions. This provides dtrictly-dffine counterexamples to theory attributing grokking solely weight norm evolution towards the Goldilocks zone 3 and also suggests new potential ways for optimizer design.
arithmetic tasks suggests that
Introduction
Modern deep learning generally relies on heavily overparameterized neural networks (NNs), which presents surprisingly strong generalization capabilities across diverse tasks 4 understanding the mechanism behind this generalizability is crucial for further advancement in the field. While model generalization performance typically improves with the training performance simultancously, the grokking phenomenon exhibits dramatically different dynamics, providing an intriguing case for the investigation of generalization. Grokking was first reported in modular arithmetic tasks where researchers used transformer model [8]. The model can rapidly reach nearly 100% training accuracy but the test accuracy remains very low, suggesting that the model only memorizes the training dataset without actually understanding them. However, after extensive further training, the model exhibits a sudden transition a state of high generalization where the test accuracy is (almost) at perfect level. This significantly delayed generalization distinguishes grokking from typical learning processes and raises fundamental challenges towards the existing generalization theory g 10
6 77. Therefore,
CCoresponding authors: entao.yang @ aarliquide.com com, Scchang7@@cityu.eddle
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Previous works have shown that the loss landscape of overparameterized NNs possesses infinitely many minima; they exhibit different generalizability and are connected with each other to form a single low-loss manifold [[]]. Therefore, many studies have attempted to propose geometric indicators that correlate with the generalizability of minima. One representative hypothesis that wider minima generalizes better than their narrower analogs [2]. However, it is difficult to measure or even rigorously define the width of a minimum in such a high- high-dimensional space. While some indicators like sharpness [13) correlates with model generalization empirically, they have also been questioned for the validity in measuring generalizability alone, due to the general scaling invariance in NNs [4]. Other indicators such as the Hessian matrix are too expensive to calculate in deep learning applications [3]. These underscore the limitations of using local geometry alone to predict generalization, which hinders the understanding of complex generalization dynamics like grokking.
Overcoming the limitations of purely geometric descriptions, a very recent work has explored applying concepts from statistical mechanics to NNs, which conceptualizes NN as physical system where parameters are the degrees of freedom and uses Boltzmann entropy calculated over parameter space to characterize the loss landscape [2]. In this context, entropy measures the logarithm of the volume of solutions in the parameter space. The work reveals a high-entropy advantage for NNs' generalizability, where the high entropy solutions (the thermodynamic equilibrium state in the context of physics) tend to exhibit better generalization than solutions obtained via classical optimizers like SGD or Adam (2). This motivates to employ the entropy framework to investigate the dynamics of grokking. We hypothesize that the transition from memorization (pre- groe-rroking) to generalization (grokking) might involve the network navigating towards the higher-entrops regions. Following the reported protocol (2). we employed the Wang-Landaal Molecular Dynamics (WLMD) algorithm to sample transformer neural networks' entropy 5(In Atest)
function of the logarithm of training loss In( Levai))) and the test accuracy Atest, for three different modular arithmetic tasks with varying complexity. We find that the grokking process can be mapped a computational glass relaxation: the training loss of a NN is equivalent to the internal energy of glass, the memorization training mimics rapid liquid cooling to form non-equilibrium glass at low temperature, and the eventual grokking emerges as a slow relaxation process towards a more stable, higher-entropy configuration that generalizes.
The paper is organized as follows: In Section 2 we review the background of generalization, grokking, Boltzmann entropy, and glass physics. In Section 3l we introduce in detail how we sample the entropy landscape of an NN with WLMD. In Section we first present entropy landscapes from three different modular arithmetic tasks, where two of them exhibits grokking while the other one is unlearnable. The results indicate that there is no entropy barrier between the memorization and generalization states. Therefore, grokking is NOT first-order phase transition, which contrasts with Rubin et al.'s previous work M. The grokking NNs also exhibit much more significant high-entropy advantages than those studied in Ref. (2). We then eliminate grokking by training with constrained weight norm 33 and find that the high-entropy advantage is significantly reduced. These findings inspire us develop toy optimizer, WanD, presented in Section [5] We demonstrate that WanD can find generalizing solutions in a competitive efficiency to AdamW while almost eliminating grokking, without any weight-norm constraint or output modifications. Finally, in Section 6 we further discuss related work and situate previous findings on grokking within our proposed framework of computational glass relaxation. We also discuss directions for future research. Code is available at http://ggthuubbcooo//taag222///
Our key contributions are as follows:
Grokking is NOT first-order phase transition. By conceptualizing neural networks as a physical system and sampling its entropy landscape, we find there is no entropy barrier between the pre-grokking memorization state and the final generalizing state, arguing against the first-order phase transition interpretation M. Instead, grokking has similar behavior to the formation and relaxation of glass, where the pre-groke gre-grokking memorization states can be attributed to the fast convergence of train loss minimization, and resembles the rapid cooling of a liquid during glass transition. The subsequent grokking process can therefore be recognized as a slow, computational glass relaxation towards a more stable, generalizing configuration. Such an analogy can describe many reported observations for grokking and inspire potential research directions for NNs, as we detailed later.
The high-entropy advantage under grokking a unique property of generalization. We find that in the entropy landscape of the grokking NNs, there is a huge gap between
2

---

the training curve and the max- entropy curve, which was previously reported as "high- entropy advantage in generalizability" for other NNs and datasets 2]. The high-entropy advantage under grokking, as we found here, is much larger. Furthermore, we show that after eliminating grokking by controlling the weight norm, the high-entropy advantage is much less prominent but still exists. This confirms that the NN without grokking does not have glassy behavior, as they are much closer to the equilibrium state, supporting the correspondence between grokking and glass relaxation. It also reinforces the understanding that higher entropy correlates with better NNs' generalizability 22 broadening its empirical support to our distinct model and task.
We design a physics-inspired toy optimizer, WanD, which finds high-norm generalizing solutions without grokking. WanD is an optimizer based on the Wang-LLndaa molecular dynamics method. In the modular addition task, WanD has similar time efficiency and generalization performance as Adam but can eliminate grokking without any explicit regularization (like weight norm constraints). Notably, WanD often finds high-norm gen- eralizing solutions, which challenges theory attributing grokking solely to weight decay [3]. This also agrees with the conclusion of Kumar et al. [[O] that regularization is not always necessary for grokking, but in a strictly-defined situation without output scaling, We argue that grokking is an issue caused by optimization dynamics and can be solved by just improving the optimizer. We believe that WanD points out a potential direction to design new optimizers which has an advantage in finding higher entropy solutions with better generalizability, and our research can encourage further work on applying well-established patistical-physice methodologies for deep learning research.
2 Related Work
Generalization has long been linked to flat minima in the loss landscape 12 In an over- parameterized NN, different minima form a connected low-loss manifold 15). Some optimizers like SGD are believed prefer flatter minima due its inherent gradient noise 16. explaining why SGD can usually yield better generalization than other optimizers [16). Therefore, many works focus on developing flatness-based measures and observe correlations with generalizability empirically
[13 [7 18 19] However, the usage of flatness as a simple predictor of generalization has been challenged by many works, mainly because it has the dimension of weight [20]. Dinh al. [14] changed the loss landscape by reconstructing the fitting function without affecting the generalizabilitg and therefore demonstrated that sharp minima can also generalize. Neyshabur et al. [21] and Feng et al. 20 discussed that flatness itself is not sufficient to ensure generalization, but needs to be combined with the weight norm to obtain suitable complexity measure. Our results support this because the combination of flatness and weight norm actually reflects the entropy, the logarithm of the parameter space volume, of the solutions with the same generalizability.
Grokking is an anomalous generalization phenomenon that was first discovered by Power et al. [8] in arithmetic tasks, describing a model's sudden transition from random- guoooss-ggssss (nearly) perfect test accuracy long after achieving perfect training performance. Different mechanisms have been proposed to explain this delayed generalization, including: (a) shifts from poor to good representation learning
[10 22 23, 24] (b) transitions from dense to sparse predictive subnetworks [25]: (c) slingshot
mechanism for adaptive optimizers [6];: (d) first-order phase transition in the network's internal
representation [I] and (e) evolution of weight norms towards the Goldilocks zone R. Our analysis
from the entropy perspective is consistent with (a), but challenges (d)(e) with well-defined
counterexamples. Detailed discussion on connections with existing work in grokking are presented in
Section 6
Boltzmann Entropy defines the number of microstates for a given macrostate in thermodynamic system. For example, entropy can be defined by the volume in phase space with certain energy value: S(U,) = kglnV/u=zi Vl= where k is Boltzmann constant and V represents the volume in phase space that correspond the internal energy Ug. Yang et al. 2 applied the concept of Boltzmann entropy to the NNs. This allowed them to improve upon the local perspective of sharpness and weight norm, and established a mapping relationship between generalization and the volume of parameter space. They found that high-entropy NN states usually generalize better than states trained with classical optimizers like SGD, which they define as the high-entropy advantage. This inspired us to use the same approach to study the relationship between entropy and generalization in grokking tasks.
3

---

Glass state is a physical state of matter, and is usually formed by rapidly cooling or compressing a liq- uid. The glass state is not in thermodynamic equilibrium. Therefore, at appropriate temperatures, the glass state will undergo structural relaxation- process where the atomic or molecular arrangement non-equilibrium glass gradually evolves toward a more stable, more equilibrated configuration over time. This relaxation process is generally regarded as kinetic process [27]. Winter and Janssen [28] compared deep NNs with structural glass and found many similarities, including the relaxation dynamics at low temperatures. As we will show in this paper, grokking is thermodynamicall similar to glass relaxation, so a promising future direction is to explore whether the conclusions [28] holds in our setup.
Methodology
In this section, we will first define the modular arithmetic task, then introduce the concept of entropy into NNs, and finally present our sampling method: the Wang-Landau Molecular Dynamics algorithm. Code available https///gttbbbooo//rhangggg//o/
3.1 Problem definition
The training part of our experiment is same as the setup of Nanda et al. (2), that is, a one-layer transformer model trained in modular arithmetic tasks. The data consists of two parts: prompts (x) (v) (p) and answer (x y). In the three tasks we study, the answers are determined by the following equations:
z0y = x x y mod p for S x. <p
x goy = x²² + y mod for S x2 < p
(I)
x xo y = x² + xy² + y mod for S x, <
Here, v, y. and p are all natural numbers. In this paper, p is dataset will be divided into a training set and a test set with a 50% fraction, by a fixed random seed. In the Appendix Figure we provide visual image of the full dataset. Transformer hyperparameters and more training details are also given in the Appendix A
prime number 67. The entire
3.2 Entropy in neural networks
There are many ways define entropy in physics and information theory. this work, we apply the definition of Boltzmann entropy S = kg in V to NNs. Here, kg is the Boltzmann constant, which we set to 1 for simplicity. physics, V is the volume of the phase space that corresponds to a particular set of observables (energy, magnetization, etc.). For an NN, we define V as the volume of the parameter space that corresponds to a particular training loss and test accuracy, and therefore S also function of these two quantities. To ensure that V is finite, we limit the parameter space by setting different reflective bounds for the parameters of different layers. These bounds are designed cover the range of parameters typically reached during normal training while avoiding being excessively large. In Sec. [42] we avoid applying these reflective bounds study the entropy advantage when grokking is suppressed. Instead, we limit our study to a hyperspherical shell of weight norm 30, a method known to effectively prevent grokking [3]. Considering grokking often occurs when the training loss is small, results presented in terms of its logarithm, ((In(Ltrai)) ), Asest). Here, Ltrain denotes the training cross cross-entropy loss, and Atest represents the smoothed test accuracy.
The entropy landscape we calculate cannot reach infinite resolution. For computational purposes, we divide (Lrran)r))) Atest) Atstt) which is 2D function, into rectangular bins, where each bin corresponds to the entropy of the NN parameters within a certain range of training loss and test accuracy. The Wang-Landau Molecular Dynamics method allows us to calculate with an arbitrary zero point. In other words, the absolute value of the entropy is meaningless, but the entropy difference between any two bins is meaningful. Therefore, we can subtract a constant from the entropy in each bin so that the minima-bin entropy is always 0.

---

3.3 Wang-Landau Molecular Dynamics (WLMD)
In this section, we will introduce the details of WLMD algorithm and its implementation in NNs. The key advantage WLMD is its ability to avoid becoming trapped [30]. In contrast to standard samplers, which tend to oversample the high-entropy regions. WLMD achieves uniform sampling across all regions of the entropy diagram upon convergence [31. This is accomplished by maintaining running estimate of the entropy landscape, which introduces an 'entropy gradient' that drives the system away from previously visited states. The convergence of the WLMD is established in [32) For a physical system made of atoms or particles of other length scales, one can use WLMD to calculate the entropy as function of the potential energy S(U). This is done by simulating particles with the following equation of motion [30]:
d4. To aS(U, 1) (U 02..... 4n) dB2
(2)
m; OU
0q:
where 91, 92,-.., Qj is the coordinate of each particle, To is the simulation temperature, m, is the mass of the particle, is the simulation time, S(U,t) is an entropy function that we update over time, and U is the potential energy of the system. Initially, S(U, = 0) O. During the simulation, we
repeatedly update SUUtt):
S(U,t + At) = t) + (w) - UD), where (t) = Fw min(t/tw l, twi/t). (3)
Here the Wang- Landau factor ft quantifies how fast we update the entropy function. In order to ensure the convergence of the WLMD algorithm, usually we make it increase linearly with time to a maximum value F and then decay inversely. 5 is Dirac function, and we approximated it with narrow Gaussian function the simulation. Uo is the potential energy of the system at this moment.
For NNs. we can set the temperature To and the mass my to 1, which does not affect the correctness of the WLMD algorithm. The particle coordinates Q1, Q2: .... In. are replaced by the NN parameters C1 C2, Cn. and the entropy we considered is S(ZWL, L.L). We choose in (Leroin) ) as the energy in broad sense TW L. and Atest the order parameter yw L. The parameter motion follows
&c LS(#WL. dt2
@S Lirasn), Atest, t) Iln (Ctrain( C2.... Cn))
dc,
Oln Leroin)
dc:
(4)
SSln Ltrain); Atests 1) Aeeest(Cr C2 ...,.... Cn)
AAtest
Oc:
8s The term Al((Le)))) JAtere is the gradient of In(train loss) over parameters. Since Atest is discrete, we
and as can be computed via numerical differentiation between adjacent bins, while Aln(Leratn) approximate its gradient using the sigmoid function, following the approach in [2]. Further details are also provided in appendix[A.3 AA.3 Reflective boundary conditions are used to implement parameter constraints for results in Sec. 4.1 and for results in Section 4.2 we restricted the parameter weight
de;
norms directly by weight rescaling (both detailed in Sec. 3.2)
The initialized velocity of the parameters follows the Boxzell-Bolzzann distribution, which is random Gaussian distribution with mean and variance k /m 1. order to maintain constant temperature during the simulation, that is, to conserve the kinetic energy of the parameters, we add the Langevin kinetic term to the velocities:
de (t + At)
dcs
de.
des (1)+ (e, cy)knT Urand di AL oj
(1) di
- -0
(5)
dl
Here, the definition of &ec is given in Equation 4 Cy is the friction coefficient, and Brand is the random velocity follows standard Gaussian distribution. In the NN parameters updating process, by continuously adding entropy to the entropy landscape, S(Zw JWL, t) will eventually converge to the true entropy distribution S(ZWL, VW L..
de
Due to space limitations, more details of WLMD and the complete algorithm are described in the Appendix A.3 To the best of our knowledge, WLMD is currently the most efficient entropy sampling algorithm, as existing alternatives are limited to models with only hundreds of parameters [2]. The computational complexity of WLMD scales linearly with the number of NN parameters, and the entropy landscape sampling can be parallelized across multiple GPUs. This enables entropy sampling for large-scale models, such transformers with hundreds of thousands of parameters.
5

---

4 Grokking: Glass Relaxation Process in Neural Network
4.1 No entropy barrier in entropy landscape
In thermodynamics, a phase transition is marked by significant change in an order parameter macroscopic, measurable quantity like density during a liquid- gas transition. By analogy, the learning dynamics of NNs exhibiting grokking appears to mimic the characteristics of phase transition: starting with rapid progress from underfitting to memorization (low test accuracy), the NN undergoes an abrupt qualitative shift to generalization (high test accuracy) after a prolonged training period. Given this, we identify the test accuracy as the order parameter for characterizing grokking. Just as a physical order parameter changes distinctly between material states, the definitive increase in test accuracy track the NN's transition into the generalization regime, making it natural and compelling macroscopic indicator of this critical learning phenomenon.
A previous study suggests that grokking is first-order phase transition process from memorization to generalization []] They used the Langevin dynamic optimizer to train two two-layer nonlinear teacher-studen models and found that the memorization and generalization states of the NN were located at two free-energy minima, thus believing that grokking a first-order phase transition that requires crossing free-energy barrier. We directly examined this statement by computing the entropy landscape as function of training loss and test accuracy to search for such a barrier. We examine three algorithmic tasks on one-layer transformer model. The three tasks range from easy to learn, to difficult to learn, and to impossible to Iearn. The training results are also given in the Appendix Figure 3 where grokking is observed for the first two tasks. Using the WLMD method, we sample the entropy landscape of training loss and test accuracy as presented in Figure | The three entropy landscapes are the results after 60M, 80M, and 20M WLMD update epochs respectively, where these durations were chosen as the points where the entropy map exhibited no substantial further change. We ran 8 WLMD processes on 4 GeForce RTX 4090 GPUs, requiring approximately 100 GPU hours per 1M epochs. Additional computational details are provided in Appendix A.2
For the first two tasks that exhibit grokking, the training trajectory is roughly "I." shaped. In the vertical part of the trajectory, i.e., from the memorization state (lower left part) to the generalization state (upper left part), we observed no entropy barrier. Since there is neither entropy barrier nor an energy (train loss) barrier, our results disproves the view that grokking requires crossing a free-energy barrier. In addition, we find that the equilibrium (maximum entropy for a given training loss) curve in the low train loss region is continuous, which also denies the first-order phase transition, as such transitions require discontinuity of the order parameter (test accuracy). For the third unlearnable task, the equilibrium state also fails to generalize, although some states present a higher test accuracy (about 20%) than random selection (about 1.5%).
(a) x+y 67
(a) x2
x 300xx"" mnd 67
1.0
25000
1.0
30000
1.0
50000
0.0
25000
50000
20000
0.5-
15000
0.6
moodbed kot
0.6
15000
40000
10000
0.4
10000
0.2
max entropy
5000
0.-
0.2
AdamW training
0.0
0.0
0.0
-10
-6
-4
-2
-15
-10
-5
-15
-10
Tattrain loss)
Inttrain loso)
Ialtrain J0s0)
Figure 1: Entropy landscape of three different arithmetic tasks (a) x + y mod 67 easy to learn, max entropy state can generalize better than training line. (b) z² + 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 + zy² + y mod 67 cannot be learned by training, max entropy state also can not generalize.
We also found huge gap between the AdamW training curve and the equilibrium state curve. We define this significant advantage of the equilibrium state over the training curve as the high-entropy advantage under grokking. This result shows that grokking actually a kinetic process rather than a phase transition in the thermodynamic sense: The NN is far from the equilibrium state during the

---

initial part of the training process, and finally reaches the equilibrium state, which generalizes well, after a long time. This process is similar glass formation and relaxation: When the liquid is cooled rapidly, the molecules do not have enough time to rearrange themselves into crystal or crystal-like structure, forming a disordered glassy state (non-equilibrium state), which then gradually evolves toward a better equilibrated state through slow structural relaxation.
4.2
Neural network without grokking does not exhibit glassy behavior
The analogy between grokking and glass relaxation can be further verified by examining the entropy landscape after eliminating grokking. According to our conjecture, just as a slowly cooled liquid will maintain thermodynamic equilibrium and will not enter the non-equilibrium glassy state, an NN without grokking will be trained along the equilibrium curve as the training loss decreases, steadily reaching the generalization state. The study by Liu et al. (3) suggested that grokking can be eliminated by limiting the weight norm of the NN parameters in an appropriate range. Following this insight, we implemented a strict weight norm constraint in our study: after each iteration step, we compute the weight norm of the neural network parameters and scale them proportionally to maintain a fixed norm of 30. The training results, presented in Appendix Figure show that grokking can be greatly eliminated under this condition. We examine three transformer models with varying widths (model dimensions). For each entropy sampling task, we use NVIDIA H100 GPUs run processes, totaling 60M epochs and consume nearly 4,000 GPU hours. Further computational details are available in Appendix A.2
(a) width 12
(h) width 64
(c) width 128
30000
1.0
40000
250000
aounthec bost Sccarss
0.8
25000
0.8
30000
0.0
200000
20000
15000
20000
0.4-
0.4
0.4
100000
10000
5000
0.2
10000
50000
AdamtW training
0.0
0.0
0.0
-10.0
-5.0
-2.5
0.0
-10.0
-7.5
-5.0
-2.5
-10.0
-7.
-5.0
-2.5
Infiruia Jese)
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 2 shows that after controlling the weight norm at 30, the Adam training curve in the entropy landscape is very close but still slightly below the equilibrium state. Moreover, under this restriction the entropy landscape cannot explore severe memorization states, which is the computational glassy state in our analogy. This result further illustrates the correspondence between glass and grokking. In addition, we note that the high-entropy advantage still exists even in the network where grokking is climinated. This agrees with the work Yang et al. (2) that high-entropy advantages are prevalent in various deep learning tasks and there is robust connection between equilibrium and generalization.
Shoows
5 WanD: Eliminating grokking without regularization
The above results show that the equilibrium state low training loss has high generalizability regardless of whether the weight norm of the network is restricted. Therefore. we believe that grokking is dynamic process caused by the optimizer's inefficient exploration of the loss landscape basins. The WLMD method we used is good at sampling high entropy states of NNs, which inspires us to design an optimizer by imitating WLMD, that can find high entropy states and therefore should also eliminate grokking. Here, we design an optimizer, WanD, based on the one-dimensional Wang-Landan Molecular Dynamics. The continuous-time dynamics of WanD as follows.
d=O(t) dt
@S(In(Liraci ((tlaaa((nn(t))))) (0(t))) Oln Öln(Lrrain (0(1)) + 77(1)
dd(()
(6)
Oin Letrin (0()))
a0(1)
dt
where (()) is a tensor of the NN parameters time I,, and In (Ttrain) the logarithm of the cross- entropy training loss. nt()) is noise given by (n,(t)n = X(2 A)kBId,,0(t t'), KBT =
7

---

is temperature term, and À is a friction coefficient. When sampling the entropy landscape using WLMD, we obtain the corresponding training loss and test accuracy the NN to locate where to add entropy. For the WanD optimizer, its training process samples the entropy landscape S(In( (tvaia)) based solely on training loss, due the fundamental constraint that can only access training data. In addition, we set up appropriate entropy bias to accelerate the NN's exploration at low training loss area. The overall WanD training procedure is summarized in Algorithm [
Algorithm WanD Optimizer
Require: Transformer model T()), the entropy bias Shiss, training data (Xtrains Utrasn and test data(Xtest> Stest)
Ensure: Transformer model I"' reach generalization: Accuracy(T((Xan Btest) Re 100%
The initial entropy: < Sbias
for each optimize epoch do
Compute the logarithm of the train loss:
In In(Lrrain) In( CCossrosEviue Entropyy (TaonTT(Xaaa)))) Ntrain));
Compute entropy gradient Egrad daw
dS(wwz)
Compute &e. atn Lerain) Eirrad dc;
dta
for each parameter c; do
Update the parameter velocity: des + des (1 cg) +
d' de
At + Langeuin
2t
dt
Update the parameter: C, + + de end for
de; At
Update the Wang- factor by Equation
3 fwL +- Fw ww
Update the entropy landscape by Equation 3 5 ((WwL) t S(ZwL) + fwwo(x LWL)
Verification:Atet Accuracy( (T(Xtest). Vtest)
end for
The training results for modular addition task using WanD or Adam are shown in Figure Both methods have comparable performance, achieving generalization after 2, 000 optimization steps. WanD has shorter running time, taking about 50 seconds per 1K steps compared to roughly 72 seconds for AdamW on a single GeForce RTX 4090 GPU. During the training process of WanD,, the weight norm of the NN parameters continued to increase, eventually reaching about 175, indicating that the NN can generalize high weight norm. This agrees with the work Kumar et al. [0O that regularization is not necessary for grokking, but in more strictly-defined context NN's output, which could be taken as effectively changing the weight norm. We present additional training results for WanD and AdamW W in the Appendix [.3 on the same task but eliminating grokking by increasing the training dataset split fraction.
do not scale
175
1.0
0.0
0.6
0.2
6.2
0
520
1000
1500
2000
300 2000
1500
2000
40
50
80
100 120 140 100 100
2000
2000
3000
optiniontiussoo steps
swr.inization stepo
wwnght FOETS
rptimootiion saype
Figure 3: Using (c) WanD and (d) AdamW to train one layer Transformer modular addition (p = 67) task. (a) Train and test accuracy over time. The grokking largely eliminated while training weight norm. These (a, b, c) together
with WanD. (b) Weight norm over time. (c) Test accuracy demonstrate that the model can generalize at high weight norm. (d) Train and test accuracy over time. Both WanD and Adam W use the same learning rate.

---

6 Conclusions and Discussion
6.1 Conclusions
In this work, we investigated grokking in the modular arithmetic tasks using a one-layer transformer, and found that Boltzmann entropy (logarithm of solution volume) as an excellent indicator of generalization. By sampling the entropy landscape of the network on training loss and test accuracy, we found that i). there is no entropy barrier between the memorization state and the generalization state, which means that the memorization state is not in an equilibrium or meta- state. Why does grokking take so long happen even without crossing free energy barriers? In the last part of this section, we postulate that the reason is the formation of glass states. ii). The equilibrium test accuracy,
function of the training loss, appears continuous. This strongly suggests that that grokking is not a first-order phase transition as argued by previous work [.. because a first- order phase transition implies discontinuity of physical quantities. Our results also indicate that the equilibrium state in the entropy landscape has a significant advantage over the training curve, which we refer as the high-entropy advantage under grokking. Such a high-entropy advantage is weakened but still apparent when grokking is eliminated via fixing weight norm. This demonstrates correlation between equilibrium state and NNs' generalization, consistent with the conclusion of Yang et al. 22. In addition, we designed a toy optimizer based on the Wang-Lander molecular dynamics method, WanD, which encourages exploration of the high entropy state for NNs' training. WanD) successfully achieves generalization with a similar time efficiency as Adam but avoids grokking. The generalization states found by WanD have a much higher weight norm compared with the one needed to eliminate grokking in traditional training [. This further supports the conclusions of Kumar et al. [10] that grokking is not solely due to weight decay as some theory proposed before 3 in a strictly-defined context.
6.2 Discussion
Based on the above conclusions, we proposed a hypothesis: grokking is similar to the process of quickly cooling a liquid to form a glass and then slowly letting it relax back to the equilibrium state. Through this analogy, we try answer two key questions about grokking:
What causes the NN to fall into the memorization state? One requirement for the formation of glass is to cool the liquid below the glass transition temperature 333 Below this temperature, the free movement of atoms is greatly restricted, thus hindering the system from evolving towards equilibrium. In grokking tasks such as modular arithmetics,, the NNs are severely overparameterized, i.e. the number of the parameters is much larger than the size of the dataset. During training, the NN can fully memorize all training data and achieve very low training loss. If we take the training loss as the energy, this is equivalent to lowering the temperature below the glass transition temperature. Another requirement for glass formation is sufficiently fast cooling rate, so that the molecules do not have time to rearrange themselves into the lower-free- free-energy crystalline structure. During the training process, the gradient descent based optimizers, especially adaptive optimizer like Adam, effectivvly implement a rapid 'cooling' process by aggressively driving parameters towards regions of reduced training loss. This rapid descent often proceeds with insufficient perturbations (caused by large learning rate, small batch size, and data accuracy, etc.) to escape from poor local minima, making NNs easier to reside in a gonegeneral loss minima and therefore form a computational glass. In contrast, the Langevin dynamics in the WanD optimizer we designed enriches the perturbations in training, thus avoiding the formation of computational glass.
What is the internal dynamics of the grokking process? can form a -ouuluuuiii glassy state that gradually approaches equilibrium through structural relaxation-a slow kinetic process (27). Due to the sluggish motion of atoms at low temperatures, the relaxation time can far exceed the timescale of the cooling process. Analogously, in grokking tasks, the equilibrium state at low training loss corresponds to good generalization. However, parameters located in the low train loss regime exhibit very small gradients, which severely limits the ability of gradient- based optimizers to effectively update the network toward this generalizing equilibrium. This dynamical slowdown results a substantial delay between the memorization and generalization phases. The effect is further exacerbated by adaptive optimizers like Adam, where the non-uniform scaling of gradients can further dampen parameter updates, increasing the likelihood of being trapped in -on-generalizable minima (34). Although the duration of the grokking process varies across
When a liquid is rapidly cooled, it
9

---

training runs, as observed in prior studies and our experiments, the transition to generalization consistently takes significantly longer than the memorization stage [8) mirroring the slow relaxation behavior characteristic of glassy systems 27
Our conclusions agree and disagree with some previous studies. There were series of work arguing that grokking is due to a transition from poor to good representation learning [0] 22 23 24]. This agrees with our analogy of relaxation from a poorly structured "glassy" (memorizing) state more ordered (generalizing) configuration. In fact, our entropy perspective offers a resolution to debates regarding solution "efficiency" and parameter norms in previous works [C 22 While Varma et al. [2]] argue that grokking requires transition to an "efficient" generalizing circuit characterized by a lower parameter norm, Kumar et al. [T0] challenges it as they identified generalizing solutions with higher norm than memorizing solutions. Our framework suggests this conflict stems from using parameter norm alone to define efficiency, while entropy is actually a better indicator since the efficient solution will have more "free" parameters and therefore a larger entropy. This is also consistent with Merrill et al. 25 who showed that memorization solution and generalization solution correspond to dense and sparse subnetworks, respectively. In the sparse subnetworks, only a few neurons participate in the memorization of the rules, and the remaining degrees of freedom contribute to the entropy, so it should have larger entropy as we observed. Kumar et al. [T]] proposed that the key determinant of grokking is the rate of feature learning, which is consistent with our glass analogy, where the rate-limited free motion of particles is important reason for the glass formation and long relaxation time. The finding by Liu et al. (4 that a fast-learning decoder can help trigger grokking also supports our analogy, as the rapid decoder acts similarly to 'astt cooling', trapping the network in glassy memorization state. Our work also supports the finding that perturbation is beneficial to the grokking learning process (26). as Wanl) is able eliminate grokking directly. Rubin al. claimed that grokking is a first-order phase transition process, but our results overturn this conclusion. This is probably because their network architecture and training data are different from ours, while ours are consistent with the original paper that discovered grokking. In addition, the high-weight norm generalization states we found through WanD also challenge the conclusions of Liu al. Due to space limitations, we provide additional discussion in the Appendix C
6.3 Limitation and future work
Limitation: i). Task scope: Our study is limited to one-layer transformers modular arithmetic tasks and does not include other grokking benchmarks (e.g., tasks in [3). ii). Interaction with explicit regularization: The relationship between "glassy dynamics" and various explicit regularization methods in grokking systems remains unclear and requires further exploration. iii). Proof-of-concept nature of WanD optimizer: WanD has only been tested in arithmetic task scenarios, is sensitive to hyperparameters, and exhibits instability near generalization.
Our hypothesis of grokking and glass relaxation invites further investigation. While focused on experiments with the original grokking tasks (8) future studies should examine whether our findings applicable for other grokking tasks. Furthermore, recent study has compared the behavior of deep neural network and structural glass with quantitative measurements (28]. and we can refer to its research method to verify whether the grokking NN is structural glass. In addition, it remains unclear how does glassy dynamics affect the effectiveness of regularization methods in removing grokking. Finally, we also aim further improve the WanD optimizer to enhance its stability and general applicability. As an algorithm based on WLMD, WanD involves multiple hyperparameters that require tuning. Our experiments identify the learning rate, temperature, and friction coefficient as the three most influential factors affecting training outcomes. We plan to further investigate empirical guidelines for effective hyperparameter selection. Techniques from statistical physies designed to accelerate equilibration, such as simulated annealing and parallel tempering, can also be applied to WanD's future refnement [35]. We anticipate that approaches such as a dynamic hybrid optimization strategy combining WanD and AdamW could help mitigate instability issues and present some preliminary results in the Appendix [B.3] Ultimately, we aim to develop WanD into optimizer applicable to broader range of machine learning tasks.
purpose
10

---

Acknowledgments
The authors thank National Natural Science Foundation of China for supporting this research (Grant 12405043, G.Z.). also thank computational resources provided by Bridges-2 at Pittsburgh Super- computing Center through ACCESS allocation CIS230096 (E.Y.)
References
[1] Noa Rubin, Inbar Seroussi, and Zohar Ringel. Grokking as first order phase transition in two layer networks. arXiv preprint arXiv:2310.03779, 2023.
[2] Entao Yang, Xiaotian Zhang, Yue Shang, and Ge Zhang. High-entropy advantage in neural networks' generalizability. arXiv preprint arXiv:2503. 13145, 2025.
[3] Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. arXiv preprint
2022.
[4] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 245. Icee, 2009.
[5] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[6] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguisties: human language technologies, volume I (long and short papers), pages 4171-4186 2019.
[7] Tom Brown, Benjamin Mann. Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 13:1877-1900 2020.
[8] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting small algorithmic datasets. arXiv preprint 221.07200..02177. 2022.
[9] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107 115, 2021.
[10] Tanishq Kumar, Blake Bordelon, Samuel J Gershman, and Cengiz Pehlevan. Grokking as the transition from lazy to rich training dynamics. arXiv preprint arXiv:2310.061 10, 2023.
[11] Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference on machine learning. pages 1309-1318. PMLR, 2018.
[12] Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Newral computation, 9(1):1442, 1997.
[13] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large- training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: 04036 2016.
[14] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference Machine Learning, pages 1019-1028. PMLR, 2017.
[15] Alexander Shevehenko and Marco Mondelli. Landscape connectivity and dropout stability of sgd solutions for over-parameterized neural networks. International Conference on Machine Learning, pages 8773-8784, PMLR, 2020.

---

[16] Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. arXiv preprint NDV:20022033995 2020.
[17] Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint 17..v::703.1100000 2017.
[18] Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where find them. arXiv preprint arXiv: 2019.
[19] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware mini- mization for efficiently improving generalization. arXiv preprint arXiv:2010.01442, 2020.
[20] Yu Feng, Wei Zhang, and Yuhai Tu. Activity-weight duality in feed-forward neural networks reveals two co-determinants for generalization. Nature Machine Intelligence, 5(8):908 -11900811,, 2023.
[21] Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generalization in deep learning. Advances neural information processing systems, 30, 2017.
[22] Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint arXiv:2309.02399. 2023.
[23] Daniel Kunin, Allan Raventós, Clémentine Dominé, Feng Chen, David Klindt, Andrew Saxe, and Surya Ganguli. Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning. Advances in Neural Information Processing Systems, 37:81157- 8 81203, 2024.
[24] Liu, Ziming and Kitouni, Ouail and Nolte, Niklas S and Michaud, Eric and Tegmark, Max and Williams, Mike. Towards understanding grokking: An effective theory of representation learning, Advances in Neural Information Processing Systems, 35:34651-34663, 2022.
[25] William Merrill, Nikolaos Tsilivis, and Aman Shukla. tale of two circuits: Grokking as competition of sparse and dense subnetworks. arXiv preprint arXiv:2303.11873, 2023.
[26] Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint arXiv:2206.0444 046771 2022.
[27] Birte Riechers, Lisa A Roed, Saeed Mehri, Trond S Ingebrigtsen, Tina Hecksher, Jeppe C Dyre, and Kristine Niss. Predicting nonlinear physical aging of glasses from equilibrium relaxation via the material time. Science advances, 8(1 :)eab19809 2022.
[28] Max Kerr Winter and Liesbeth MC Janssen. Glassy dynamics in deep neural networks: A structural comparison. Physical Review Research, 7(2)022010, 2025.
[29] Neel Nanda, Lawrence Chan, Tom Licberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05211 2023.
[30] Christoph Junghans, Danny Perez, and Thomas Vogel. Molecular dynamics in the multi- canonical ensemble: Equivalence of wang-landau sampling, statistical temperature molecular dynamics, and metadynamics. Journal of chemical theory and computation, 1843):184-1187, 2014.
[31] Kim, Jaegil and Straub, John E and Keyes, Thomas. crristicl--emppratuure Monte Carlo and molecular dynamies algorithms. Physical review letters, 97(5): 050601, 2006.
[32] Fort, Gersende and Jourdain, Benjamin and Kuhn, Estelle and Lelièvre, Tony and Stoltz, Gabriel. Convergence of the Wang-Landau algorithm. Mathematics of Computation, 2015.
[33] Jean-Pierre Hansen and lan Ranald McDonald. Theory of simple liquids: with applications to soft matter. Academic press, 2013.
12

---

[34] Zhou, Pan and Feng, Jiashi and Ma, Chao and Xiong, Caiming and Hoi, Steven Chu Hong and others. Towards theoretically understanding why sgd generalizes better than adam in deep learning. Advances in Neural Information Processing Systems, 33, 21285-21296, 2020.
[35] Daan Frenkel and Berend Smit. Understanding molecular simulation: from algorithms to applications. Elsevier, 2023.
[36] Chaudhari, Pratik and Choromanska, Anna and Soatto, Stefano and LeCun, Yann and Baldassi, Carlo and Borgs, Christian and Chayes, Jennifer and Sagun, Levent and Zecchina, Riccardo. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We listed and described claims made in this work both in the abstract and the introduction. The main claims are made based on our experimental results (section 4.
and
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer:
[Yes]
Justification:
We discussed the limitations of our work and future directions Section d
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiitiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: We propose computational glass relaxation hypothesis based on the observed entropy landscape in grokking tasks and demonstrate the validity of the analogy. We also verify this result using an optimizer we designed, WanD.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We give the sampling method in the Section and provide the hyperparameters required for our work in the Appendix. Readers can reproduce the results based on the code we provide.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
15

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We provide code about i). Sampling entropy using WLMD ii). The optimizer WanD we designed.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
httppp.ccppubcccguuuceeee Codeeubmii ssi criionPooiccy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide experimental details the Appendix. Full details are given in the code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Due the computational cost of entropy landscape sampling, we fixed the same data split and random seed for both training and sampling procedures, so error analysis is not applicable.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoo bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We give the hardware environment used in the experiment and the computa- tional cost in the Appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't iidn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://///eeuaipp.cccccppbb////a
Answer: [Yes]
Justification: Our work of grokking conform with the NeurIPS Code of Ethics in every respect.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work towards understanding generalization of grokking is of academic interest only and does not have direct societal impact.
Guidelines:
17

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work on grokking in modular arithmetic tasks do not include such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
We give proper credit to the assets used mainly via citations and textual descriptions.
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
18

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms use in the datasets.paperssiii paperswi thcode. www.coo//aaaaa datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We currently do not introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
19

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
We do not involve LLMs in this work.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should not be described.
20

---

A Additional Experimental Details
A.1 Model and Datasets
We followed the work of Nanda et al. (2 for the model setup of modular arithmetic tasks. We use the same one-layer transformer model for the programs of AdamW optimizer training, WanD) optimizer training. and WLMD entropy sampling. This transformer model has 4 attention heads by default, each with width of 32, yielding model dimension of 1 128, and attached with MLP layer with width of 512. In the experiment of eliminating grokking by controlling the weight norm, we modified the attention head width (8/16/32) and the MLP layer width (128/256/512). computational cost of WLMD entropy sampling, we selected a smaller prime number 67 as the modulus for our task. Although some works 3 [9] choose 113 as the modulus, there still exist grokking in arithmetic tasks with 67 as the modulus, so we think this setting is appropriate. We also use the cros-yy loss as our loss function. For more details on the model's hyperparameters. please refer to the provided code.
order reduce the
A.2 Training Details and Cost
For modular arithmetic tasks training by AdamW, we use a full batch size on GeForce RTX 4090 GPU, training for 10000 optimization steps with weight decay 1 or fixed weight norm 30. The hyper- parameters for training with AdamW are given in Table [ For toy optimizer WanD training, since it is not yet fully developed, there are many hyperparameters that directly affect the generalization of the training results. Among them, the learning rate, temperature, and friction coefficient the three key factors that determine the training results. We provide a set of hyperparameter combinations in the code that can generalized well.
Table I: AdamW training hyperparameters
Operation
Model Dimension
MLP Width
Learning Rate
Weight Decay
Weight Norm
x + y
32
128
1e-3
0
Fixed at 30
y
64
256
1e-3
0
Fixed at 30
I + y
128
512
le-3
0
Fixed at 30
y
128
512
3e-3
I
Not fixed
x² +
128
512
5e-3
1 Not fixed
x² + xy
+
128
512
le-2
1 Not fixed
For WLMD sampling, we use 4 GPUs to run 8 processes. Each process uses the same data split and updates the entropy landscape according to a different random seed. A single process runs 100,000 epochs at time, then calculates the average entropy of all processes and updates the entropy landscape. When running WLMD processes on 4 GeForce RTX 4090 GPUs, takes about day for each process to update 1M epochs. When running 8 WLMD processes on 4 NVIDIA H100 GPUs, it takes about 16 to 17 hours for each process to update 1M epochs. In order for the entropy landscape to converge to its true distribution, tens to hundreds of millions of epochs are typically required. For detailed hyperparameter settings, please refer to the code we provide.
A.3 WLMD details
Smoothed accuracy:
To ensure well-defined gradient of test accuracy with respect to NN parameters in WLMDD, compute smoothed test. accuracy using a sigmoid function, following 2 The detailed procedure for obtaining smoothed accuracy is outlined in Algorithm 2 In our work, we choose the smoothed accuracy slope s =
When sampling the entropy landscape of grokking NNs, we are interested in the regime with lower training loss, which corresponds to the memorization state and the generalization state. However, areas with high training loss usually have higher entropy (as evidenced by randomly initialized networks that generally possess high training loss), which will waste a lot of time exploring in the early stages of sampling. In order to prevent the simulation program from wasting time outside the range of our interest, we set additional entropy bias:
Entropy bias:
if In (n(Lerain) Lo, Soiae (In( ((In(Lrrain) Lo)?: else (bine (In( = 0. (7)
21

---

Algorithm 2 Smoothed Accuracy
Require:
Transformer model T()), dataset (X {x1, 22, I2, In Y {ūr, D2 ... Dn smoothed
accuracy slope 6
Ensure: The value of smoothed accuracy close to the true accuracy, and the gradient of smoothed accuracy with respect to parameters has a finite value.
for each (Iz, (s) in (X, Y) do
Get the highest logit among all wrong output classes: (1 Vi_urong) where Vi_wrong
Y = T(x:)
Get the logit of the correct output class: Pı(y 9:) Compute the smoothed accuracy:A; = Siggmoid (p:ly = 9:) -pe(y = yi Vi_wrong)) .8)
end for
Compute the total smoothed accuracy: = Σ A,/n
If the training loss is outside the range interest (In(Lerain) > parameters, this entropy bias will provide additional entropy gradients to to force the NN to move towards lower training loss.
Lo) when updating the NN
The overall WLMD entropy sampling procedure is outlined in Algorithm 3 entropy loss for our modular arithmetic task and logarithmic scale of the training loss is used to better examine the entropy landscape near zero training loss.
We employ the cross-
Algorithm WLMD Entropy Sampling Method
Require: Transformer model T(.). its current entropy landscape S, the entropy bias Soiae. training
data (Xtrnnn Ytrein Jand test (ata(Xtestt Mteat);
Ensure: Convergent entropy landscape S(Iw 9WL) where IWL (r((Lrrin)) D Atesti
Initialize: S + Sbias;
for each sampling epoch do
Compute the logarithm the train loss and smoothed test accuracy:
In(Irrain) In( CerCCrossEsnrppp Loss(T(X (T(Xrain)) Mtvain));
Atest = Smoothed Accuracy ( (T (eest) (rest);
Compute entropy gradient Egrad 8S( (.(WL)) and Vyrad 05(zwL.9wZ),
L
0gwz
The combined loss combined -W S.W.FFgre gradd I.ggaddd
Do Leombined backward() and get de: for each parameter c; do
dec
Update the parameter velocity by Equation de, Update the parameter: C2 +- + de end for
5
do: dt
do: dt (1 ce) +
Pe, At + ULangevini do
A
Update the Wang-Landau factor by Equation Update the entropy landscape by Equation TWL: y ywL);
3. fWL +
(IWL, (WwL:WwL +- S(FWL.WEL) + fwr5(x
end for S S Suiasi
B Additional Results
B.1
Visualization of modular arithmetic tasks
First, we show visualization of the three modular arithmetic tasks in the Figure 4 dataset with 67*67=4489 data points can be viewed as 67*67 pixel image, where the x-axis and y-axis coordinates of the pixels correspond to the two numbers (x) and (v) involved in the operation, and the color of the pixel corresponds to the answer (x y). Half of the pixels are randomly selected as the training set, and the other half as the test set. Training the transformer model to learn the modular arithmetic operation is equivalent to completing the entire image according to the regularity of the image after being masked. In particular, the numbers the arithmetic operations are just equivalent to different characters, and they can be replaced by any other different characters without
A complete
22

---

affecting the validity of the modular arithmetic task. In order to observe the image intuitively, we dye the pixels with light dark blue according the corresponding numbers from small to large. is not difficult observe that the visualization images of simpler tasks have more obvious regularities. For the task that cannot be learned, its visualization image is close to random noise.
(a) x mod 67
(b) 2+y mod 67
(G) x x"3+xy 24y mod 67
20
20
20
40
40
40
40
50
10
20
60
Figure 4: Visualization of three different arithmetic tasks. (a) x y mod 67. (b) x² + mod 67. (c) x -xy² + mod 67. Our protocol is equivalent to randomly taking 50% of the pixels from the image as training set and trying to complete rest. The more complex the expression, the more visually irregular the image is.
B.2 A training results of modular arithmetic tasks
The AdamW training results of these three modular arithmetic tasks shown in Figure 5 The figure shows that as the expression becomes more and more complex, the learning difficulty of the NN also increases, from easy learn, to difficult to learn, impossible learn. The generalization delay is more significant in the difficult task, which means grokking is more pronounced.
x mod
x2+y mod
(o) (x)x mod
1.0
0.8
test
0.8
0.8
acturary
0.4
0.4
0.4
0.2
0.2
0.0
10"
10'
10'
10a
10a
10'
10'
10'
10a
10'
10'
10"
10'
optimization staps
oytimization stops
optmization stags
Figure 5: Training results of three different arithmetic tasks: (a) x + + mod 67 easy to learn, but the hysteresis of the test accuracy shows grokking. (b) x² + mod 67 is hard to learn, the grokking more obvious. (c) xyz + y mod 67 cannot learned.
We also train the transformer in the modular addition task with constrained weight norm. The weight norm restricted to the Goldilocks zone 33, and we set to 30 for three NNs with different widths. We tested three NNs with different widths mainly because Yang et al. [2] observed that the high-entropy advantage becomes more significant with smaller width. The training results are shown in Figure 6 In fact, these three NNs reach generalization after almost the same number of optimization steps, and none of them exhibits grokking. Combined with the entropy landscape in Figure the main text, believe that the relationship between high-entropy advantage and NN width is not obvious the grokking task.
B.3 Training results of AdamW and WanD under different dataset splits
Increasing the training dataset split fraction is an effective way to eliminate grokking [8) We tested the training results of AdamW and WanD on the y mod 67 (Figure 7 and x² + +y mod 67 (Figure
23

---

(a) width 32
(b) wicth
(6) width 128
1.0
0.0
0.6
0.6
0.4
0.6
0.2
0.2
0.2
0.0
0.0
0.0
staps
optimizatioo steps
optimicatioo steps
Figure 6: Training results of modular addition (p 67) task with three transformer models that have different widths. Here weight norm is constrained to fixed value of 30.
8) tasks the same learning rate 0.003, with the training dataset fraction being 50%, 60%, 70%, and 80%, respectively. Here we anneal WanD with AdamW so that the update weight of WanD increases first and then decreases during training, thus avoiding the instability problem at the end of WanD training. Specifically, in each iteration, the NN parameters are updated by the WanD optimizer with weight a and by the AdamW optimizer with weight a, where a (2Atest - 1)2.
(a)
(b)
(c)
(d)
1.4
10-
00-
6/-
203
1289
183
200
31D
300
200
203
200
420
(h)
4.0
1.0
Optinimntten
Figure 7: AdamW and WanD training results of x + mod 67 task with different training dataset fraction. Top row:AdamW: bottom row: WanD; (a)()) 50%, (b)(f) 60%, (c)(g) 70%, and (d)(h) 80% training dataset split fraction. The hyperparameter of WanD: kT=0.5, IritiionCoefficient-0... Lorrrannnoosooodddy---o-so
The results show that WanD has a significant training efficiency advantage in tasks with severe grokking, and even in tasks where grokking is greatly eliminated, WanD can still achieve generaliza- tion ahead of AdamW. This result demonstrates the great potential of WanD as a toy optimizer and motivates us continue improving it.
C Additional Discussion
Robustness of entropy landscape diagrams: In this work, we visualize the entropy landscape with respect to the logarithm of the training cross-entropy loss and the smoothed test accuracy. Due to the symmetry between training and test datasets, modifying the training loss formulation is conceptually equivalent to altering the test loss. Our reported results are based the smoothed

---

(b)
(c)
(d)
62-
010
2230
1000
6602 s000
19900
300
1900 2500 1900 3000
300
1300
1090
209
830
2003
Oytidaatio: Sags
19li.tauinn Stage
Oyttniataa Supe
(c)
(f)
g
(h)
1.0
0.-
24
0.2
300 1230 1990 7800
3500
1299 1899 2000 2000 3000
500
1299
1822
1200
Figure 8: AdamW and WanD training results of 22 + mod 67 task with different training dataset fraction. (d) (e)(fXg))t): (alce>//bb training dataset split fraction. The hyperparameter of WanD: kT=0.7, frictionCoeffccientee 6, maxTrainl Loss Loss three times the efficiency of AdamW in (a).
00%/60%/70%///
psssTottudy=
worth noting that WanD in (e) achieves generalization with nearly
test accuracy, whose inverse can be interpreted as form of loss function. We have also conducted internal experiments using cross-entropy test loss in place the smoothed accuracy. The resulting entropy diagram consistently support the same conclusion: entropy barrier is not observed, and only single connected high-entropy region exists, indicating no sign of first order phase transition. This consistency confirms that our findings are robust to the specific choice of the loss function.
Sparse subnetworks: As shown in [25]. grokking is associated with sparse subnetworks in which only a subset of neurons is involved rule memorization. From an entropy perspective, parameters outside these functional subnetworks do not affect generalization. Perturbing such parameters preserves model performance, implying that the corresponding solution set occupies a larger volume parameter space. Since entropy defined as the logarithm of this volume, neural networks with sparse subnetworks exhibit higher entropy.
WanD is optimizer built on the one-dimensional WLMD framework. It excels at identifying high-entropy states, particularly the generalized solution, via molecular dynamics, leveraging the fact that generalized solutions possess higher entropy low training loss. Prior to our work, there is another entropy-guided optimizers named Entropy-SGD 36 have been proposed. Below, we briefly highlight the theoretical distinctions between WanD and Entropy-SGri SGD:
Other entropy-based optimizer:
Different notions of entropy. The core distinction lies in the entropy definition: WanD targets the global Boltzmann entropy, reflecting the total volume the solution space, whereas Entropy-SGD focuses on local entropy, which captures the flatness of individual minima.
Different methodologies. WanD estimates entropy via ID WLMD sampling, while Entropy- SGD relies on Langevin dynamics to approximate local entropy gradients.
Optimizers like Adam often progress slowly due to limited gradient signals during the memorization- to- on-generalization transition [10]. In principle, Entropy-SGii could potentially accelerate grokking if its identified high local-entropy states exhibit strong generalization. rigorous empirical comparison between these approaches remains valuable direction for future work.
25


Reviews:
[{"review_id": "kb5WpWYEyh", "reviewer": "Reviewer_b947", "summary": "This paper presents a novel interpretation of the grokking phenomenon in neural networks through the lens of statistical physics, particularly computational glass relaxation. It employs the Wang-Landau Molecular Dynamics algorithm to estimate entropy landscapes in modular arithmetic tasks and demonstrates that grokking does not correspond to a first-order phase transition, contradicting prior work. Additionally, the authors propose a physics-inspired optimizer, WanD, capable of eliminating grokking without explicit regularization.", "questions": "Implementation and Stability of WanD Optimizer: The WanD optimizer appears promising, but the manuscript lacks details on its computational overhead, sensitivity to hyperparameters, and stability. Could you provide more information about:\n1. The additional computational cost compared to AdamW and If WanD is shown to be efficient, stable, and generalizable with minimal tuning, that would strengthen the practical contribution and could positively affect the score.", "limitations": "I am not very familiar with this area and lack some of the necessary background in physics, which made the paper challenging for me to fully understand.", "rating": 4, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. Novel Theoretical Perspective: The paper provides a fresh and compelling physical analogy for an AI phenomenon. Viewing a neural network as a physical system and grokking as analogous to glass transition/relaxation is an innovative framing. This perspective synthesizes ideas from statistical mechanics with deep learning, potentially opening new avenues for understanding generalization. The analogy is well-grounded in the results: the authors show that pre-grokking fast training is analogous to rapidly quenching a liquid into a glass (a non-equilibrium state), and the grokking phase is akin to a slow structural relaxation toward equilibrium.\n\n2. Entropy Landscape Analysis: Methodologically, the use of the Wang-Landau Molecular Dynamics algorithm to sample the entropy landscape of neural network parameters is a significant strength. This approach goes beyond conventional training-loss landscapes by explicitly estimating the Boltzmann entropy (volume of solutions) at different performance levels. The paper thus moves beyond local geometry (sharpness or flatness of minima) to a global thermodynamic view of the parameter space.\n\n3. Connection to Prior Work: The authors do a good job situating their contribution relative to prior explanations of grokking. They explicitly discuss how their entropy-based view can account for earlier observations.  At the same time, they challenge two prominent prior hypotheses – the phase transition interpretation and the weight-norm-driven explanation.\n\nWeaknesses:\n1. Limited Scope of Experiments: A notable limitation is that all experiments are conducted on relatively small-scale, synthetic tasks (modular arithmetic problems with a small transformer). \n\n2. Computational Complexity and Feasibility: The approach of sampling an entropy landscape with WLMD, while powerful, is computationally intensive and may not be practical beyond toy settings. The paper does not clearly report the computational cost of obtaining the entropy landscape for even these small models.\n\n3. Clarity for Non-Physicists: The paper assumes a certain comfort with physics terminology which might pose a barrier to some readers.\n\n4. Missing or Understated Insights: The paper touches on how its findings relate to earlier hypotheses (like representation learning or the slingshot effect), but these connections remain somewhat abstract. For example, it is asserted that the entropy perspective can explain observations of representation shift or sparse subnetworks, yet the reader must infer the mechanism (presumably, high-entropy solutions correspond to learned sparse structures or improved features). These connections would be more convincing if explicitly demonstrated or at least elaborated.", "was_revised": true, "final_justification": "I am not familiar with the related work in the area, but I feel the proposed approach can have a significant impact. Based on the improvement discussed in the rebuttal, I keep my initial recommendation.", "rebuttal": "We thank the reviewer for their positive assessment and insightful suggestions, which will help us strengthen the manuscript.\n\n$\\textbf{Weakness}$:\n\n$\\textbf{W1. Limited Scope of Experiments}$. While we acknowledge the reviewer's point, our one-layer transformer model (with different width) and three modular arithmetic tasks is $\\textbf{the most representative experiment}$ setup for grokking research (see Ref [3, 8, 10, 25, 28] in our manuscript as examples). Focusing on this ‘benchmark of grokking’ allows us to better compare it with more prior work. \n\nWe also agree with the reviewer that expanding our analysis to a wider range of grokking-like tasks is an important direction for future work. One such example is using the same task and model design reported in Ref [1], which concludes grokking is a first order transition but uses a slightly different definition of the grokking phenomenon, focusing on the 'sharp drop in test loss' (see Figure 5 in Ref [1]).\n\n$\\textbf{W2. Computational Complexity and Feasibility.}$ We thank the reviewer for this important point. We reported the computational cost in Appendix A2 together with the details of our experimental configuration but recognize that we did not highlight it clearly in the main text. We will update section 3 in the final version to clearly present it.\n\nWhile the computational cost to obtain the converged entropy landscape of this transformer model (with more than $\\textbf{400 thousand parameters}$) is relatively high, we believe this analysis itself represents a significant advancement. To the best of our knowledge, the WLMD method is the $\\textbf{most efficient}$ entropy sampling algorithm available now, because other methods are limited to models with only hundreds of parameters (as explained by Ref [2]). Besides, considering it takes us ~10,000 NVIDIA 4090 GPU hours to obtain the converged entropy landscape shown in Figure 2 with parallel computation, and that the computational complexity scales linearly with the number of parameters, it is also computationally feasible to scale this work to larger models with tens of millions of parameters if more advanced hardware is available. \n\n$\\textbf{W3. Clarity for Non-Physicists.}$ We agree with the reviewer that our manuscript uses some physics terminology, like thermodynamic equilibrium, glass relaxation, and Boltzmann entropy. This is mainly  due to its nature as a statistical physics inspired work. We defined and explained these concepts when we mentioned them and will also utilize the additional page allowance in the final version to add more intuitive analogies to ensure the work is self-contained for a broad audience. We also welcome the reviewer to suggest specific concepts that may need additional clarification and we will supplement them in the final version. \n\n$\\textbf{W4. Missing or Understated Insights.}$ We thank the reviewer for highlighting the connection to prior work as a strength of our study. Due to the length limitations, we only have a section to briefly discuss these connections at the end of the manuscript. We will add an additional appendix to discuss more details in the final version. The reviewer’s reasoning about the underlying mechanism of the connection between high entropy solution and learned sparse structure / representation shift is also precise. Therefore, below we detailed it as an example of the explanation we will add.\n\nRef [24] has shown that grokking is related to sparse subnetworks, where only a few neurons participate in the memorization of the rules. From our entropy perspective, for the NN with this generalized sparse subnetwork, since the parameters outside the sparse subnetwork do not contribute to the generalizability, their perturbation would not harm generalization. Therefore, they constitute a solution set with a larger volume in the parameter space. The entropy of this solution set is defined as the logarithm of the volume in the parameter space. Thus, the entropy of the NN with sparse subnetworks is larger. \n\n\n$\\textbf{Questions}$:\n\n$\\textbf{Q1.}$ We are encouraged by the reviewer's interest in WanD and would like to provide more experimental details about it below.\n \n$\\textbf{- Efficiency}$: Due to the scope of this work, we only tested WanD on the modular addition task. Compared with AdamW with the same learning rate (3e-3 for results presented in Figure 3a and 3d), they need a similar number of optimization steps (\\~2000) to reach generalization. We have also run another experiment with a learning rate of 7e-3 and got similar results (both generalize at \\~1400 steps). The running time for WanD is even shorter, which is \\~50 seconds per 1K steps (\\~72 seconds for AdamW).\n\n$\\textbf{- Sensitivity to hyperparameters}$: currently we consider WanD as a $\\textbf{toy optimizer}$ that demonstrates how a physics-inspired optimizer can successfully navigate the grokking landscape. As a WLMD-based algorithm, we acknowledge that there are many hyperparameters that need to be set. Our experiments find that the learning rate, temperature, and friction coefficient are the $\\textbf{three key factors}$ that determine the training results. \n\n$\\textbf{- Stability}$: as the reviewer may notice, the current optimization trajectory of WanD is relatively unstable, especially in the later stage of training. We are working on improving WanD right now and have obtained encouraging preliminary results. For example, the stability issue of WanD can be addressed by employing a $\\textbf{dynamical hybrid optimization}$ between WanD and AdamW, where WanD dominates in the early stage and AdamW contributes more in the end. This enables us to use WanD to escape the initial memorization phase and AdamW to stabilize final convergence. We find it can achieve the same level of stability as AdamW on the modular addition task, with an improved efficiency and no delayed generalization. We will report detailed results in future work soon. Therefore, we are optimistic that WanD could eventually become an optimizer applicable to general machine learning tasks."}, {"review_id": "I6qA6hEFmI", "reviewer": "Reviewer_xfSb", "summary": "This work draws an analogy between\ngrokking in neural networks (a sudden drop in test error \nwell after the network has memorized the training data) \nand the behavior of particles in a liquid after\na sudden temperature drop. \nBuilding on this connection, the authors construct a phase\ndiagram of the entropy (volume of weight configurations) of a one-layer transformer with a given training loss and test accuracy. \nThey observe that:\n1) in the grokking regime (low training loss), the entropy\nlandscape does not exhibit any barrier, challenging the \nbelief that grokking is a first-order phase transition;\n2) the solutions that maximize the entropy \nfor a fixed training loss generalize well and do not \nexhibit grokking.\n\nBased on these insights, the authors introduce a new \nalgorithm that targets high-entropy solutions and avoids grokking without \nthe need to control weight norms.", "questions": "1. I found the intuition behind the notion of\n    ''conical degree'' in Section 6 somewhat confusing.\n    Is there a specific reason to believe that the cutoff\n    induced by fixing the training loss admits a simple\n    geometric interpretation?\n2. Figure 1.(b) seems to be the key experimental\n    result, as this is where grokking is most apparent for AdamW. The max-entropy curve in this panel\n    appears to rise sharply towards perfect test accuracy. How did you ensure that the continuity of the \n    max-entropy curve is not a finite-size effect?\n3. Several algorithms have been proposed to explicitly target high-entropy solutions in the traditional loss landscape sense (for example https://iopscience.iop.org/article/10.1088/1742-5468/ab39d9). How do those methods compare empirically to the one proposed in this paper, both in terms of generalization performance and existence of grokking?\n4. In the sampling of the entropy landscape using molecular dynamics (Equation 4), the training loss and test error appear to play symmetric roles, especially in the absence of regularization in the training loss. This makes me wonder, how robust are the resulting phase diagrams to changes in the choice of training loss?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n1. This is a very interesting contribution to the\n    ongoing discussion around the origin of grokking.\n2. Overall, the paper is cleanly written, the story\n    is clear and the experimental results are well presented.\n\n## Weaknesses\n\n1. The presentation of the two core algorithms\n    (for the simulation of the entropy landscape and targeting high-entropy solutions)\n    lack details, both in terms of how and why they\n    work, and how they were implemented. It would\n    be helpful to give more intuition on\n    them, especially for a broader\n    ML audience, which is unlikely to be familiar with\n    Wang-Landau molecular dynamics.\n2. My understanding is that \n    the proposed WandD algorithm is computationally expensive. A discussion about its complexity/convergence rate \n    would strengthen the work.", "was_revised": true, "final_justification": "The current work offers an interesting perspective on the grokking phenomenon and brings value to this line of research. All questions raised during the discussion have been addressed. I maintain my score.", "rebuttal": "We thank the reviewer for the valuable comments and for recognizing our work's novelty and the paper's clarity. Below we provide our response on a point-by-point basis, to address the reviewer's concerns and improve our manuscript.\n\n$\\textbf{Weakness}$\n\n$\\textbf{W1. The presentations of WLMD and WanD lacks details}$. We agree that a more intuitive explanation of our core algorithms would be helpful for a broad audience and will add the following clarifications, including new algorithm boxes, to the final version of paper (section 3 and 5 respectively). Our attached code can also serve as a reference if needed. \n\n$\\textbf{WLMD}$: We explain the implementation of WLMD in section 3 and provide the reference of WLMD (Ref [29]) in the main text, which contains the detailed formula derivation. We will incorporate the following explanation to it in the final version.\n\nThe key idea behind WLMD is it can avoid getting trapped. Unlike standard samplers that oversample the high entropy regions, it can achieve a uniform sample across all regions of the entropy diagram after convergence. This is realized by maintaining a running estimate of the entropy landscape, which adds an ‘entropy gradient’ to push it from the already-visited states. \n\n$\\textbf{WanD}$: Equation 7 gives the core dynamics of WanD, and we will add the following explanation around it in the final version.\n\nThe intuition behind WanD is to utilize the entropy landscape to train models more effectively. It employs 1D WLMD to sample the entropy of training loss and promote exploration of high-entropy regions. This allows WanD to navigate away from the low entropy memorization state, and therefore eliminate the grokking process and finally find high-entropy solutions. We also have an entropy bias (defined in Equation 6) in WanD, which encourages exploration of the low-loss regions.\n\n$\\textbf{W2. WanD's Complexity}$. We thank the reviewer for their interest in our toy optimizer, WanD. We reported its computational expense in Appendix A2 together with the details of our experimental configuration but recognize that we did not highlight it clearly in the main text. We will update section 3 in the final version to present it. \nThe WanD optimizer is actually quite efficient based on our test on the modular addition task. It needs a similar number of optimization steps (\\~2000) to reach generalization, compared with AdamW with the same learning rate (3e-3 for results presented in Figure 3a and 3d). We have also run another experiment with a learning rate of 7e-3 and got similar results (both generalize at \\~1400 steps). The running time for WanD is even shorter, which is ~50 seconds per 1K steps (\\~72 seconds for AdamW).\n\n$\\textbf{Questions}$\n\n$\\textbf{Q1}$. We thank the reviewer for this question, and we agree that the \"conical degree\" concept was not sufficiently explained and could be confusing. Our intention was to provide a high-level geometric intuition of the generalized solution. Our goal was to suggest why Boltzmann entropy (solution volume) is a more comprehensive measure for generalization, compared with the local flatness alone. \n\nConsidering that the ‘conical degree’ is not a formal result of our paper, we will remove it from the final version and only discuss that entropy is a more comprehensive measure since it reflects volume.\n\n$\\textbf{Q2}$.  We thank the reviewer for this important question. The continuity of the max-entropy curve in Figure 1 is supported by the primary signature of a continuous transition: the absence of a low entropy barrier. The entropy landscape we finally get is smooth and only has one connected high-entropy region bridging the memorization and generalization states.\n\nIn case the reviewer is asking about the finite size effect of neural networks studied, we would like to clarify the scope of our work. Our objective is to study the finite-width networks that are used in practice, rather than infer properties of the infinite-width networks. Therefore, our conclusions cannot be extrapolated to the theoretical infinite-width limit.\n\n$\\textbf{Q3}$. We appreciate this insightful question. While a direct empirical comparison is a quite interesting area for future work, we can offer a conceptual comparison between our approach and optimizers like Entropy-SGD, which also leverage the idea of entropy.\n\nThe key distinctions are:\n\n$\\textbf{- Different entropy}$: The primary difference lies in the definition of entropy. WanD targets the global Boltzmann entropy, which measures the total volume of the solution space. In contrast, Entropy-SGD targets local entropy, a measure of the flatness of a specific minimum. \n\n$\\textbf{- Different methodologies}$: WanD computes entropy via 1D WLMD sampling, whereas Entropy-SGD uses Langevin dynamics to estimate the local entropy gradients.  \n\n$\\textbf{Regarding grokking}$: Optimizers like Adam progress slowly due to the insufficient gradient motivation for memorization-to-generalization transitions (Ref [10]). In principle, Entropy-SGD may accelerate grokking if its identified high local-entropy states can generalize. We believe a rigorous empirical comparison remains an excellent direction for future research.\n\n$\\textbf{Q4}$. This is an excellent question about the robustness of our findings. As the reviewer observed, symmetry exists between training and test datasets. Therefore, a change in the formula of the training loss is equivalent to a change of the test loss. \n\nThe results in our manuscript use the \"smoothed test accuracy\": the inverse of which can be considered as a type of loss function. We have internally also experimented with the cross-entropy test loss, instead of the smoothed accuracy. The resulting entropy diagram leads to similar conclusions: there is no entropy barrier and only one connected high-entropy region, providing no evidence of a first order phase transition. This suggests our conclusions are robust given the specific choice of loss function. We will add a note about this to the Appendix in the final version."}, {"review_id": "uENaNIDpXR", "reviewer": "Reviewer_WLHH", "summary": "- Authors compare grokking to a slow relaxation of a many-body system into an equilibrium state after a rapid cooling process\n- they argue against a first-order state transition interpretation, as brought forward by previous work, by investigating the boltzmann entropy landscape as a function of train loss and test accuracy and finding no discontinuity in the order parameter during the transition process\n- authors argue for a “high entropy advantage” in the grokking scenario. The training trajectory in the (train loss, test acc) landscape differs vastly from a more beneficial trajectory, the maximum entropy curve, which achieves high test accuracy much quicker. they claim a broad correlation between high (boltzmann) entropy solutions and generalizability in neural networks.", "questions": "- What the relationship between $U$ and $(L_{train}, A_{test})$ ?\n- I don’t really understand the WL simulation. How exactly is $\\frac{\\partial S}{\\partial L_{train}}$ calculated? You say via numerical differentiation of adjacent bins, but that estimate is very bad at the beginning of the simulation right? Don’t you need to sample a lot of parameter configurations to estimate that? Maybe a proof of how it converges to the right estimate would be helpful.\n- Also a little algorithm box explaining the exact steps to simulate and update NN parameters would be helpful.\n- Why do you need to conserve kinetic energy? Otherwise updates become too large as potential energy is transferred into kinetic?\n- you conclude that grokking is not a first order phase transition because you don’t see a discontinuity in the order parameter you choose (test accuracy). Does that mean that no order parameter that one could choose would be discontinuous?\n- why does the need to cross a free-energy barrier imply a first order phase transition? because crossing wouldn’t happen unless one jumps right to the other side, implying a discontinuity (like quantum tunneling)?\n- How would a discontinuity present itself in Fig 1? Would the yellow line be disconnected?\n    - If that is the case, could one not always reduce the size of the update (i.e. learning rate) to make it continuous? Or would grokking just never appear with such a low learning rate in that view?\n- Why can Fig 1(c) not be learned by training?\n- How much of the training time in Fig 1(b) goes into the vertical ascent of the yellow curve vs the horizontal move along the train loss?\n- how does the conical degree combine weight norm and flatness? it would help to expand a little on that, it’s a little weird to introduce this notion in the conclusion.\n- i don’t understand why the glassy interpretation implies that generalization takes a long time.\n- do you actually check the gradient norms during the relaxation process to corroborate your hypothesis?\n\nNits:\n\n- L177 define A_test and L_train\n- L178: calculates → calculate\n- I was confused by the term “relaxation” until the very end (L354), expecting some relaxation of constraints as is common in optimization. maybe explain in the beginning that you call it relaxation because of the term “structural relaxation”.", "limitations": "consider citing and discussing the relationship to https://arxiv.org/abs/2205.10343v2\n\noverall I like the view this paper takes, but there are some holes in explanation and experimental and theoretical rigor. I am willing to adjust my score if my weaknesses and questions are addressed to my satisfaction.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 4, "strengths_and_weaknesses": "Strengths\n\n- quite clean experiments with nice interpretations\n- interesting hypothesis and theory\n- a refreshing perspective on grokking\n\nWeaknesses\n\n- the binning in (train loss, test acc) seems somewhat arbitrary. I understand that it is necessary for the simulation and the calculation of $S$, but I can imagine that different binning schemes can have vastly different outcomes if the landscape is not very regular. How do you account for that?\n- the impact of this work is not super clear to me. while it is an interesting hypothesis and seems to have some confirming evidence, what should I take away from this for the general community? the entropy to generalization link is certainly broadly interesting, but that contribution belongs to [2]. I am not convinced about the utility of the WanD optimizer.\n- there should be a proof (or at least a link to one) that the WLMD sampling converges to the true $S$ at the places where it is sampled. There are so many possible parameter configurations that it sounds non-trivial to get a grasp of true entropy by iterating this algorithm.\n- the authors claim to overthrow the first-order transition hypothesis from Rubin  et al [1] and attribute the differences to data and training setup. It seems weird that they don’t attempt the same setup as [1] to see whether they reach the same conclusion there, since they claim the disproving of [1] as a main contribution.", "was_revised": true, "final_justification": "my questions have been answered.", "rebuttal": "We thank the reviewer for the encouraging comments on our work, especially on the originality of our physics-inspired idea. We also greatly appreciate the reviewer’s effort in sharing such a thorough and insightful list of questions, which helps us to further polish the manuscript. Due to the max length limit, we have provided concise responses for some points below, but we look forward to discussing any points in more detail during the discussion period.\n\n$\\textbf{Weakness}$\n\n$\\textbf{W1. Choice of binning size}$: Our choice of binning is a balance between computational cost and the resolution of the entropy landscape. According to previous numerical analysis for Wang-Landau methods ([Belardinelli et al., 2008, Analysis), the bin size determines the final resolution but does not alter the converged landscape qualitatively. We use the log scale for train loss to focus resolution on the low loss regime. \n\n$\\textbf{W2. Main takeaway}$: Our primary contribution is a new physics-inspired mechanism for grokking, which goes beyond the general entropy-generalization correlation in Ref [2]:\n- We show that grokking can be understood as a computational glass relaxation: a slow relaxation from a kinetically trapped, low-entropy, memorization state to a high-entropy generalization state.\n- We provide counterexamples to the first-order phase transition and the weight norm determined hypotheses for grokking.\n- Our WanD optimizer serves as a practical demonstration of our theory.\n\n$\\textbf{W3. WLMD convergence proof}$: The guaranteed convergence is a key reason why we choose the Wang-Landau method. We cited the reference of WLMD (Ref [29]) in the main text, which contains the detailed formula derivation process. We will add [Fort et al. 2015 Convergence] for mathematical proof.\n\n$\\textbf{W4. Different set up from Ref [1]}$: We chose to use the same setup as Ref [28] because it is the most standard one for grokking experiments (e.g. Refs [3, 10, 25]), focusing on $\\textbf{delayed generalization after overfitting}$. In contrast, the setup and definition of grokking in Ref [1] is slightly different from the original grokking discovered in Ref [8]. It focuses more on the ‘sharp drop in test loss’ . Therefore, we believe the setup from Ref [28] is a more appropriate choice for our investigation. We agree that applying our analysis to other grokking-like phenomena, such as the one in Ref [1], is an important direction for future work.\n\n$\\textbf{Questions}$\n\n$\\textbf{Q1}$. In our analogy, the training loss ($L_{train}$) maps to the potential energy (U), as both are minimized during optimization. The test accuracy ($A_{test}$) does not map to energy. Instead, it serves as a macroscopic $\\textbf{order parameter}$ measuring generalization.\n\n$\\textbf{Q2}$. Yes, the term $\\partial S / \\partial L_{train}$ is calculated via numerical differentiation of adjacent bins (see SL_grad() in our code). The reviewer is correct that the estimate is poor initially and this is a well-understood aspect of the Wang-Landau method. The algorithm is designed to automatically fix it over time by using the inaccurate estimates of the entropy to bias the sampling, forcing the simulation to explore the under-sampled regions. This guarantees the eventual convergence to the true entropy distribution (Ref [29] and [Fort et al., 2015]).\n\n$\\textbf{Q3}$. We thank the reviewer and will add algorithm boxes for both WLMD and WanD. \n\n$\\textbf{Q4}$. While the reviewer's intuition is correct that controlling the kinetic energy stabilizes the simulation, the primary reason for its conservation is to satisfy the $\\textbf{theoretical requirements}$ of the WLMD method. Since the system's temperature is directly proportional to its average kinetic energy and our goal is to sample the canonical ensemble at a constant temperature (WLMD convergence requirement), we must conserve it.\n\n$\\textbf{Q5}$. This is an excellent question and we thank the reviewer for it. Our primary evidence for concluding that grokking is not a first-order phase transition is the $\\textbf{absence of a free-energy barrier}$ (see details in our response to Q6). The continuous nature of our chosen order parameter serves as complementary evidence. \n\nWe believe the test accuracy is a good choice of order parameters because they are the most appropriate one to describe the macroscopic phenomenon of grokking. While it is theoretically possible to manually construct a discontinuous order parameter for networks during grokking, it would be an ill-suited one that fails to describe the process clearly. For example, we can define a binary parameter called ‘perfect generalization’, which is 1 only at 100% test accuracy and 0 otherwise. Such an order parameter is $\\textbf{discontinuous by definition}$, but it is insufficient to describe how generalization evolves during grokking. \n\n$\\textbf{Q6}$. Yes, the reviewer’s reasoning is correct on a high level. In physics, a first-order transition is defined by a $\\textbf{discontinuity in the first derivative of the free energy}$, which is the ‘jump’ mentioned by the reviewer. \nThe need to cross a free energy barrier implies first order transition because it indicates the existence of a $\\textbf{thermodynamically unfavorable intermediate state}$ (like a nucleation seed in water freezing). This prevents a smooth, continuous transformation and forces the system to change discontinuously. \n\n$\\textbf{Q7}$. The key is to distinguish between the $\\textbf{training trajectory}$ (the yellow line) and the underlying $\\textbf{equilibrium state}$ (the max-entropy curve). A first order phase transition would not necessarily present as a disconnected training trajectory. Instead, its defining feature would be a $\\textbf{discontinuous jump in the equilibrium max-entropy curve}$. We would observe the memorization and generalization phases as two distinct, high-entropy regions separated by a low-entropy barrier. \n\nIf grokking was a discontinuous first order transition, reducing the step size (learning rate) would not make it continuous. A smaller LR only shrinks the steps that model takes along training trajectory, but cannot erase the gap if it exists. \n\nFinally, a low LR does not necessarily prevent grokking. In fact, previous work has found that larger learning rates can sometimes help eliminate grokking [Lv et al., 2024, Language]. This is also consistent with our glass relaxation analogy, because a larger LR is more likely to provide the system with enough \"energy\" to escape from the low-entropy memorization state.\n\n$\\textbf{Q8}$. This task is unlearnable because the training data is insufficient to uniquely determine the underlying rule and is also consistent with with the original work on grokking (Ref [8]).\n\n$\\textbf{Q9}$. The training curve in Figure 1b is corresponding to Figure 5b in the Appendix. The initial horizontal movement (memorization) is rapid and takes ~50 epochs. The subsequent vertical ascent (grokking) is much slower and takes ~2000 epochs. \n\n$\\textbf{Q10}$. We thank the reviewer for this feedback and we agree that introducing the ‘conical degree’ concept in the conclusion may lead to potential confusion. Our intention was to provide a brief geometric intuition for why Boltzmann entropy (solution volume) is a more comprehensive measure for generalization than local flatness. As it is not a formal result, we will $\\textbf{remove it}$ from the final version to improve clarity.\n\n$\\textbf{Q11}$. This is a great question that gets to the core of our analogy, the separation of timescales. For a physical glass, the initial rapid cooling is a fast process, but the subsequent structural relaxation from the initial glassy state to the stable ideal glass state is extremely slow. In grokking, the initial training is like the cooling to a memorization state with near-zero training loss. The network is then kinetically trapped, as the small gradient provides little ‘force’ for large updates. The network can only find the generalization state (ideal glass) via the cumulation of many small stochastic updates, which takes a long time.\n\n$\\textbf{Q12}$.  Yes, we checked the gradient norm during the AdamW training and the results support our hypothesis. During the initial memorization phase, the gradient norm is large (>1e-1). However, once the training loss is minimized, the gradient norm drops to an extremely small value (\\~1e-7) and remains small during grokking. The drop in gradient norm is the mechanism of the kinetically trapped state, because it leaves the optimizer with little signal to make large updates.\n\n$\\textbf{Nits}$. We will edit L177 and L178 in our final version. We also agree with the reviewer that the term ‘relaxation’ may confuse readers unfamiliar with its meaning in physics and will add clarification to section 2.\n\nWe also thank the reviewer for the highly relevant suggestion and will cite [Liu et al. 2022 Towards] in the final version. Their finding that a fast-learning decoder can help trigger grokking supports our analogy, as the fast decoder acts as the ‘fast cooling’ that traps the network in a glassy memorization state."}, {"review_id": "hNcq3Jed0s", "reviewer": "Reviewer_zyTv", "summary": "This work investigates, from a theoretical point of view, the grokking phase in NNs. The analysis is recast as a physical system modelled by computational glass relaxation, where the parameters are degrees of freedom, and energy is the corresponding training loss. This approach enables the authors to associate the memorization process with the rapid cooling of a liquid into a non-equilibrium state and the subsequent generalization with the slow relaxation towards stability. In this setup, one can define the NN Bolzman entropy. Experiments are conducted using transformers on modulo arithmetic tasks and reveal the connections between high-entropy minima and generalization ability. A novel toy optimizer that can avoid grokking is presented as a practical outcome of the previous analysis.", "questions": "I list some minor concerns:\n- (pag. 4) Is there any limitation in constraining the NN parameters to the range $3 / \\sqrt m$? \n- (pag. 4) It is stated that in Sec. 4.2  the weight norm is limited to 30. How is this enforced during training without impacting the learning dynamic?\n- In Sect. 5, WanD is tested only on the modular addition $p=67$ task that exhibits the grokking phase. Nonetheless, it would be important to add at least one experiment without grokking to ensure that WanD can be used in a general training setup when the nature of the task to be solved is unknown in advance.", "limitations": "The limitation section is missing. I think it is important to add a discussion on the current limitations of the study worked out on toy modular arithmetic tasks with one-layer transformer.", "rating": 5, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "The paper introduces a compelling and original perspective by modeling grokking as a computational analog of glass relaxation. This physical analogy helps explain the delayed generalization phase in neural networks, using entropy landscape analysis to challenge the prevailing interpretation of grokking as a first-order phase transition. The proposed optimizer (WanD) envisions the potential practical impact of the study. More experiments with Wand could be added to support the theoretical findings.", "was_revised": true, "final_justification": "I am not familiar with the related work in the area, but I feel the proposed approach can have a significant impact. Based on the improvement discussed in the rebuttal, I keep my initial recommendation.", "rebuttal": "We sincerely thank the reviewer for their positive feedback and for recognizing the originality and significance of our physics perspective on the grokking phenomena. Below we respond to the three questions raised by the reviewer.\n\n$\\textbf{Q1}$:  This constraint does theoretically impose boundaries on the transformer. As no prior studies establish parameter ranges for such models, we derived layer-wise bounds empirically by observing post-AdamW training parameter distributions. Therefore, we believe it would not qualitatively change our results. We will add the methodological details about the constraint determination in the revised manuscript.  \n\n$\\textbf{Q2}$:  For AdamW, we adopt Liu et al. 's protocol in Ref [3]: after parameter updates by the AdamW optimizer, weights are linearly scaled to enforce the condition that the norm is 30. WLMD sampling implements the identical scaling procedure (see the provided grokking_WLMD.py as a reference). \n\n$\\textbf{Q3}$: We thank the reviewer for this constructive suggestion. For this paper, our primary goal was to introduce WanD as a $\\textbf{proof-of-concept}$ to demonstrate that an optimizer explicitly targeting high-entropy states can successfully eliminate grokking by itself. \n\nThe reviewer is absolutely right that the crucial next step is to evaluate WanD's performance on standard, non-grokking tasks to ensure its general applicability. While this is out of the scope of this manuscript, we are actively working on it right now. As a first step towards a more robust optimizer, we have developed a $\\textbf{dynamical hybrid WanD-AdamW}$ approach, where WanD dominates in the early stage and AdamW contributes more in the end. This enables us to use WanD to escape the initial memorization phase and AdamW to stabilize final convergence. We believe this is a promising step to make WanD a generally applicable optimizer and will report its performance on a wider range of tasks in future work.\n\n$\\textbf{Limitations Supplement}$: \n\nWe thank the reviewer for pointing out the absence of a dedicated limitations section. Due to the length limit, we only briefly cover it in the last paragraph of our manuscript.  We will expand it in the final version with the following points:\n\n- Scope of the tasks: one-layer transformers in modular arithmetic tasks, excluding other grokking tasks (e.g., tasks in Refs [1, 3, 10]).\n\n- The proof-of-concept nature of WanD optimizer: future work is needed to improve its stability and test its robustness in general machine learning tasks.\n\n- Interaction between the 'glassy dynamics' and different explicit regularization methods in grokking systems."}]

Author rebuttal:

Meta-review: The work offers a physics-inspired perspective on grokking, viewing it as computational glass relaxation and thereby challenging explanations based on first-order phase transitions and weight norms. The approach is demonstrated through a simple grokking-resistant optimizer.


Reviewers praised the work’s originality, clear framing, and clean experiments, viewing it as a solid, thought-provoking contribution that advances understanding of grokking and inspires new optimizer designs.
Decision: Accept (spotlight)