Grokfast: Accelerated Grokking by Amplifying Slow Gradients — part 2/4

Cleaned-text code points [15998, 31500). Bibliography omitted here; source.md includes it.

Although the grokking phenomenon was first reported in the algorithmic dataset, Omnigrok [ Liu et al., 2022b ] shows that such behavior can also be observed in a diverse set of tasks with larger and more complex datasets. This section validates the efficacy of our accelerating algorithm, Grokfast , for those various tasks and models that exhibit the grokking phenomenon.

### 4.1 Algorithmic Data

We first train the same task with the same model as in Section 2 using our new Algorithm 2 . This is the same modular multiplication task devised to report the grokking phenomenon [ Power et al., 2022 ] . Consuming much smaller computational resources ( × 50 \times 50 less) compared to Algorithm 1 , exponential moving average effectively captures the slow variation of the gradients necessary for accelerating the delayed generalization. Under the grokking phenomenon, the validation loss of the model first increases before it decreases again later during the late generalization stage as depicted in Figure (baseline). This is well aligned with our state transition model interpretation in Q2 of Section 2.3 . High difference in the validation loss implies that the generalization route B → C \mathrm{B}\rightarrow\mathrm{C} is much longer than the overfitting route A → B \mathrm{A}\rightarrow\mathrm{B} in the parameter space. In contrast, models under our Grokfast training algorithms show significantly smaller peak in the validation loss as shown in Figures and . This implies that Grokfast effectively keeps the generalization route B → C \mathrm{B}\rightarrow\mathrm{C} close to the global optimum at state C \mathrm{C}\, . We will revisit these conjectures in Section 5 with more visualization.

We also conduct ablation studies to find out the effect of hyperparameters λ \lambda\, , α \alpha\, , and weight decay for our Grokfast algorithm with an EMA filter. The optimal hyperparameters are found with grid search as in Section 2 . Figures through summarizes the results. Recalling that our main idea is at the design of a low-pass filter, the momentum parameter α \alpha of Algorithm 2 as well as the window size parameter w w of Algorithm 1 are equivalent to the cutoff frequency of the underlying filter. Experiments in Figures through as well as those in Figures 4 and 7 show that there exists a sweet spot in cutoff frequency that corresponds to the generalization-inducing gradient signal. From our empirical studies, we recommend λ ∈ [ 0.1 , 5 ] \lambda\in[0.1,5] and α ∈ [ 0.8 , 0.99 ] \alpha\in[0.8,0.99]\, . The weight decay is, like in typical optimization problems, dependent on the task of interest.

### 4.2 MNIST

Besides the simple algorithmic reasoning task, where the data is relatively simple, Liu et al. [2022b] report the similar delayed generalization can also be observed in many conventional tasks if the model goes through a special treatment. To demonstrate the generalizability of our Grokfast modification of the optimization process, we try to accelerate the speed of generalization under those reported models and tasks. The first is a three-layer ReLU-MLP trained for MNIST classification task [ Deng, 2012 ] which exhibits the grokking phenomenon. Figure 10 summarizes the results, showing that our Algorithm 2 successfully accelerate the delayed generalization. With α = 0.8 \alpha=0.8\, , λ = 0.1 \lambda=0.1\, , and wd = 2.0 \mathrm{wd}=2.0\, , the delay until grokking is reduced by × 22.0 \times 22.0\, . Moreover, the final evaluation accuracy becomes higher from 89.8% to 91.2%.

### 4.3 QM9

In the next experiment, we train a graph convolutional neural network (GCNN) trained for a molecule dataset QM9 [ Ruddigkeit et al., 2012 , Ramakrishnan et al., 2014 ] . Since this task does not have an accuracy measure to compare the speed of convergence, we instead compare the convergence speed of the validation loss. With the same setup as in Omnigrok [ Liu et al., 2022b ] , elaborated in Appendix B , we apply Algorithm 2 with α = 0.9 \alpha=0.9\, , λ = 1.0 \lambda=1.0\, , and wd = 0.01 \mathrm{wd}=0.01 to obtain the results in Figure 10 . The validation loss drops faster and by a larger margin under Grokfast .

### 4.4 IMDb

Finally, we train a 2-layer LSTM [ Hochreiter and Schmidhuber, 1997 ] network for sentiment analysis in the IMDb dataset [ Maas et al., 2011 ] under the grokking phenomenon [ Liu et al., 2022b ] . Figure 11 compares the baseline with the model trained with an optimizer modified by Algorithm 2 with α = 0.98 \alpha=0.98\, , λ = 2.0 \lambda=2.0\, , and wd = 10.0 \mathrm{wd}=10.0 . We visualize the convergence speed and quantitatively compare the best validation loss/accuracy. This experiment section suggests that Grokfast generally boosts performance and convergence speed in diverse tasks under the grokking phenomenon.

## 5 More Discussion

### 5.1 Difference between Algorithm 2 and the Momentum in Typical Optimizers

The lines 7-8 of Algorithm 2 take a similar form to the momentum variable, which is frequently used in optimizers in deep learning frameworks. However, notable differences exist: (1) Instead of using the scaled momentum as a parameter update, we use the smoothened gradient as a residual , which is added to the gradient before it is fed into the optimizer. Rather, the formula is more similar to Nesterov’s momentum; however, the filtering is applied before the optimizer, which is different from typical applications of Nesterov’s momentum such as NAdam [ Dozat, 2016 ] . (2) The line 7-8 is applied to the gradients independently to the underlying optimizer. The optimizer can be of any type unless it is of the first-order gradient descent-based. Low-pass filtering the gradients g ⁡ ( t ) g(t) has the same effect as filtering the post-optimizer parameter updates u ⁡ ( t ) u(t) as mathematically explained in Appendix A with SGD and variants, and empirically proved in the previous sections with Adam [ Kingma and Ba, 2014 ] and AdamW [ Loshchilov and Hutter, 2018 ] optimizers.

### 5.2 Visualizing Trajectories

In this final section, we elaborate on our state transition interpretation of grokking introduced in Section 2.3 . Our signal space model of the training dynamics allows us to interpret the training of a model as a random drift of the state in the parameter space. To visualize the dynamics, we collect all the 423k parameters of the Transformer decoder [ Vaswani et al., 2017 ] used in the experiment in Figure 2 for all iterations, and conduct the PCA to obtain the most salient projection of the parameter space. The sequence of evolving models are projected onto the space as in Figure .

Regarding state transition interpretation of grokking, we can observe the followings: First, Figure 12 suggests that, in the baseline setup, the model drifts through a significantly longer pathway from the overfitting (state B \mathrm{B}\, , 500 steps) to its full generalization (state C \mathrm{C}\, , 300k steps), compared to the initial state (state A \mathrm{A}\, , 0 steps) to the overfitting state (state B \mathrm{B} ). However, under Grokfast , the ratio between the two distances AB ¯ \overline{\mathrm{AB}} and BC ¯ \overline{\mathrm{BC}} in the parameter space becomes more even. This is further acknowledged by Figure and Table 2 showing the distances between the models at each state. Moreover, Table 2 suggests that the distances AC ¯ \overline{\mathrm{AC}} between the initial and the final state becomes much ( × 16 \times 16 ) shorter with our Grokfast algorithm. Although the generalization accuracy, the training accuracy, the training loss, and the validation loss at the final state (state C \mathrm{C} ) are similar in both the baseline and Grokfast as showcased in Figure 2 , we cannot simply say that the states C \mathrm{C} of baseline and of Grokfast belong to the same network state. Likewise the state B \mathrm{B} of the baseline and of Grokfast are different. Figure shows average deviation of parameter weights from the initialization point during training of the model under grokking phenomenon. Interestingly, at achieving overfitting at state B \mathrm{B}\, , the model under our algorithm deviates × 8 \times 8 further from the initial point than the baseline does, with × 5 \times 5 smaller standard deviation in distances from the initial state A \mathrm{A}\, . This suggests that although both algorithms exhibit overfitted behavior at state B \mathrm{B}\, , intermediate model instances at these states form distinct set of parameters with possibly different topologies. These observations support our interpretation to regard the grokking phenomenon as a state transition between at least three distinct states. The role of Grokfast is then to provide supervision towards an alternative optimum much nearer from the initial points than the baseline optimizer.

Lastly, the model trained with our Grokfast algorithm shows hundredfold smaller variances of the distances than the baseline as claimed in Table 2 . This implies that training under Grokfast algorithm is much more deterministic than under typical first-order optimizers. This is possibly related to the similarity between the low-pass filtered gradients from small minibatches with normal gradients from larger minibatches. However, we have also demonstrated in Section 2.3 that using only the slow, more deterministic component of gradients and completely neglecting the original gradients lead to instability. Therefore, further investigation is needed to find out the source and the role of this determinism from our Grokfast algorithm, and the reason of its benefits when jointly applied with the faster, more stochastic gradients from baseline optimizers.

## 6 Related Work

#### Grokking.

The recently discovered grokking phenomenon [ Power et al., 2022 ] signifies one possibility of overparameterized neural networks generalizing (and reasoning) beyond memorization of the given dataset. Most of the works, thereafter, focus on identifying its mechanism. Recent studies associated grokking with a double descent phenomenon in the DNN mapping geometry training dynamics [ Humayun et al., 2023 ] , the speed of pattern learning [ Davies et al., 2023 ] , and the sizes of models and datasets [ Huang et al., 2024 ] , wherein the validation error initially increases and then decreases with the expansion of model parameters [ Nakkiran et al., 2021 , Belkin et al., 2019 ] . To investigate the internal roles of each model component during grokking, Nanda et al. [2023] employed mechanistic interpretability, a technique from the XAI domain, and revealed that grokking may not occur abruptly but rather exhibits an internal progress measure. Their assertion posits that the model captures slow, generalizing patterns, underscoring the critical role of proper optimization. Interestingly, while weight decay amplifies the double descent effect [ Pezeshki et al., 2022 ] , it contributes to enhanced generalization in grokking scenarios [ Power et al., 2022 ] . Liu et al. [2022b] found more examples of grokking in various tasks and analyzed their training mechanism through loss landscape analysis. Thilak et al. [2022] found a similarity between grokking and the slingshot mechanism in adaptive optimizers. Barak et al. [2022] argued that optimizers reach delayed generalization by amplifying sparse solutions through hidden progress. Regularizers such as weight decay [ Nanda et al., 2023 ] and the choice of the optimizer [ Liu et al., 2022a ] are highlighted as important factors in training a model that groks. Our work is a continuation of this discussion by providing a generalizable tool for the practical study of the grokking phenomenon. Through our discussion, we suggest a state transition model of the grokking and visualize the trajectory of the model weights in the parameter space during training.

#### Optimization techniques.

At the core of the study of grokking lies optimization techniques [ Thilak et al., 2022 ] . Studies have shown that generalization patterns of the model vary significantly depending on various optimization methods [ Power et al., 2022 , Gromov, 2023 ] . Power et al. [2022] demonstrated that various factors related to optimization, such as (mini-)batch training [ Li et al., 2014 ] , the choice of optimizer, weight decay [ Loshchilov and Hutter, 2018 ] , noise injection [ Zur et al., 2009 ] , dropout [ Srivastava et al., 2014 ] , and learning rate, influence the model’s grokking pattern. Nanda et al. [2023] argued that grokking does not occur without proper regularization. Further, they demonstrated that techniques such as weight decay, L2 norm, and dropout induce grokking, but L1 norm does not. On the other hand, Thilak et al. [2022] argued that grokking can occur without explicit regularization, attributing this to the optimizer’s “visible slingshot mechanism” acting as an implicit regularizer. Liu et al. [2022a] suggested using a larger learning rate for the input embedding segment, facilitating unified learning of the generalization pattern. Unlike these revisiting of the known training techniques, we started from a state space model and a dual domain of the training dynamics. This led us to develop an optimizer augmentation algorithm, Grokfast , that can be applied to any existing first-order optimizers to accelerate the grokking effect for practical usage.

## 7 Conclusion

Our reinterpretation of the deviation of each model parameter into a random signal over training iteration allows us to separate gradient updates into fast-varying and slow-varying components. By amplifying the latter with low-pass filtering, we can bring forward the moment of sudden late generalization, i.e., grokking, reducing the number of required training iterations by up to × 50 \times 50\, . Our comprehensive experiments and analyses suggest that our state space interpretation and the frequency representation of the training dynamics is useful for studying the grokking phenomenon.

## Appendix A Frequency Responses of the Parameter Updates under Grokfast

In Section 1 , we have silently assumed that under a first-order optimizer u ⁡ ( g ⁡ ( t ) , t ) u(g(t),t)\, , amplifying the low-frequency components of the gradient signal g ⁡ ( t ) g(t) of an arbitrary parameter θ ⁡ ( t ) \theta(t) over a discrete timestep t t has the same effect of amplifying the low-frequency component of the parameter updates u ⁡ ( t ) = u ⁡ ( g ⁡ ( t ) , t ) u(t)=u(g(t),t)\, . This section mathematically elaborates on the effect of gradient filters h ⁡ ( t ) h(t) to the parameter update signals u ⁡ ( t ) u(t) in the most frequently-used type of optimizers: SGD with momentum.

Stochastic gradient descent with optional momentum term is the simplest and the most widely used optimization algorithm in the deep learning communities. Here, the parameter update u ⁡ ( t ) = θ ⁡ ( t + 1 ) − θ ⁡ ( t ) u(t)=\theta(t+1)-\theta(t) of a parameter θ ⁡ ( t ) \theta(t) at timestep t t and its intermediate momentum m ⁡ ( t ) m(t) is defined as: m ⁡ ( t ) \displaystyle m(t) = μ ​ m ​ ( t − 1 ) + ( 1 − τ ) ​ g ​ ( t ) , \displaystyle=\mu m(t-1)+(1-\tau)g(t)\,, (8) u ⁡ ( t ) \displaystyle u(t) = − η ​ m ​ ( t ) , \displaystyle=-\eta m(t)\,, (9) where μ \mu is the scalar momentum, τ \tau is the dampening constant for the momentum, and η \eta is the learning rate. This class of optimizers can be thought of as linear systems with state m ⁡ ( t ) m(t) that receives an input g ⁡ ( t ) g(t) to produce an output u ⁡ ( t ) u(t)\, .