Grokfast: Accelerated Grokking by Amplifying Slow Gradients — part 1/4
Cleaned-text code points [0, 15998). Bibliography omitted here; source.md includes it.
# Grokfast : Accelerated Grokking by Amplifying Slow Gradients
###### Abstract
One puzzling artifact in machine learning dubbed grokking is where delayed generalization is achieved tenfolds of iterations after near perfect overfitting to the training data. Focusing on the long delay itself on behalf of machine learning practitioners, our goal is to accelerate generalization of a model under grokking phenomenon. By regarding a series of gradients of a parameter over training iterations as a random signal over time, we can spectrally decompose the parameter trajectories under gradient descent into two components: the fast-varying, overfitting-yielding component and the slow-varying, generalization-inducing component. This analysis allows us to accelerate the grokking phenomenon more than × 50 \times 50 with only a few lines of code that amplifies the slow-varying components of gradients. The experiments show that our algorithm applies to diverse tasks involving images, languages, and graphs, enabling practical availability of this peculiar artifact of sudden generalization. Our code is available at https://github.com/ironjr/grokfast .
## 1 Introduction
Grokking is a recently discovered phenomenon where generalization is achieved long after a model overfits to the training data. The phenomenon was first reported by Power et al. [2022] for a two-layer Transformer [ Vaswani et al., 2017 ] trained using a simple algorithmic dataset. Later, Liu et al. [2022b] have shown that similar artifacts are observed for various model architectures trained with a variety of datasets, including images, languages, and graphs. Many theory-oriented works have tried to justify the effect by relating the grokking phenomenon to the previously known double descent phenomenon [ Davies et al., 2023 , Huang et al., 2024 ] , yet its cause and sufficient conditions have not been fully characterized.
Apart from theoretical studies, this work takes a practitioner’s standpoint to take advantage of the grokking phenomenon. In previous reports on grokking [ Power et al., 2022 , Liu et al., 2022b ] , generalization happens only after more than a tenfold of training iterations after overfitting. This high demand of computational resources means less practical appeal to general machine learning practitioners who are often under dire resource constraints. Achieving faster generalization in those overfitting systems is, therefore, a necessary step to fully explore the potential of this unusual behavior. Our goal is, to this end, to accelerate the grokking phenomenon.
In the example training curve of a model under grokking in Figure 2 , the dynamics of the validation loss is few orders of magnitude slower than the dynamics of the training loss. The change in the losses is a direct consequence of the change in the parameter values throughout the training session. Hence, Figure 2 suggests that the parameter updates under grokking takes effect in two different timescales: the fast-varying component of the parameter updates contributes to the rapid overfitting, and the slow-varying component contributes to the slow generalization.
We begin by treating the change of value u ( t ) u(t) of each model parameter θ \theta over training iteration t t from an optimizer as a discrete (random) signal over time. As the optimizer iterates the training data, the value of each parameter θ ( t ) \theta(t)\, , the loss l ( t ) l(t)\, , and its gradient g ( t ) ≔ ∂ l ( t ) / ∂ θ ( t ) g(t)\coloneqq\partial l(t)/\partial\theta(t) drift with respect to the guidance from a sequence of randomly selected mini-batches sampled at each iteration t t\, : θ ( t + 1 ) = θ ( t ) + u ( g ( t ) , t ) = θ ( 0 ) + ∑ τ = 0 t u ( g ( τ ) , τ ) . \theta(t+1)=\theta(t)+u(g(t),t)=\theta({0})+\sum_{\tau=0}^{t}u(g(\tau),\tau)\,. (1) The parameter update function u ( t ) = u ( g ( t ) , t ) = θ ( t + 1 ) − θ ( t ) u(t)=u(g(t),t)=\theta(t+1)-\theta(t) provides a simple abstraction of the underlying optimizer. Different instances of iterative optimizers including SGD with various hyperparameters, e.g., learning rate and momentum, can be dealt with this notation.
Treating the optimization process as a collection of discrete random signals u ( t ) u(t) allows us to consider its dual representation U ( ω ) U(\omega) in the frequency domain . Taking the discrete-time Fourier transform ℱ \mathcal{F} of u ( t ) u(t) with respect to the training iteration t t\, , we obtain the spectral representation of the sequence of changes of a specific parameter θ \theta\, : U ( ω ) = ℱ { u ( t ) } = ∑ t = 0 T u ( t ) e − i ω t , U(\omega)=\mathcal{F}\{u(t)\}=\sum_{t=0}^{T}u(t)e^{-i\omega t}\,, (2) where T T is the total number of training iterations in this specific training session. Under our assumption that the delayed generalization of grokking is a consequence of the slow-varying component of the parameter updates u ( t ) u(t)\, , the grokking phenomenon is directly related to the low-frequency part of the dual representation U ( ω ) U(\omega)\, . Further, for a first-order optimizer u ( g ( t ) , t ) u(g(t),t)\, , an almost unanimous choice in deep learning applications, the gradients g ( t ) g(t) are linearly correlated with the parameter updates u ( t ) u(t)\, . Therefore, we can relate the low-frequency part of the gradient signal G ( ω ) = ∑ t = 0 T g ( t ) e − i ω t G(\omega)=\sum_{t=0}^{T}g(t)e^{-i\omega t} with the slow generalization under grokking. Our hypothesis is that amplifying this low-frequency component of G ( ω ) G(\omega) accelerates the speed of generalization under the grokking phenomenon .
In the following sections, we empirically demonstrate this hypothesis with a simple low-frequency gradient amplifier in various scenarios. These include tasks involving various network architectures including Transformers [ Vaswani et al., 2017 ] , MLPs, RNNs and (Graph-)ConvNets and diverse datasets such as algorithmic data, images, languages, and graphs that are treated to exhibit the grokking phenomenon [ Liu et al., 2022b ] . Our method is simple, taking only a few lines of additional code and is applicable to most of machine learning frameworks such as PyTorch [ Paszke et al., 2019 ] , with × 50 \times 50 faster exhibition of grokking as shown in Figure 2 .
## 2 Amplifying the Low-Frequencies of the Stochastic Gradients
### 2.1 Filter Design
Amplifying the low-frequencies of the gradients g ( t ) g(t) can be achieved by adding a low-pass filtered signal g ( t ) g(t) to itself. Let h ( t ) h(t) be a discrete-time low-pass filter (LPF) defined over the training iteration t t\, . For simplicity, we assume a univariate time-invariant low-pass filter h ( t ) h(t) uniformly applied across every model parameter θ \theta\, . Using a convolution operator ∗ *\, , we denote the modified gradient g ^ ( t ) \hat{g}(t) as: g ^ ( t ) = g ( t ) + h ( t ) ∗ g ( t ) , \hat{g}(t)=g(t)+h(t)*g(t)\,, (3) which can then be plugged into the parameter update function u u of the optimizer: u ^ ( t ) = u ( g ^ ( t ) , t ) = u ( g ( t ) + h ( t ) ∗ g ( t ) , t ) . \hat{u}(t)=u(\hat{g}(t),t)=u(g(t)+h(t)*g(t),t)\,. (4) In the dual domain, equation ( 3 ) is equivalent to: G ^ ( ω ) = G ( ω ) + H ( ω ) G ( ω ) = ( 1 + H ( ω ) ) G ( ω ) , \displaystyle\hat{G}(\omega)=G(\omega)+H(\omega)G(\omega)=(1+H(\omega))G(\omega)\,, (5) where H ( ω ) = ∑ t = 0 T h ( t ) e − i ω t H(\omega)=\sum_{t=0}^{T}h(t)e^{-i\omega t} is the transfer function of the filter h ( t ) h(t)\, . Our goal can therefore be restated as to design a filter h ( t ) h(t) with low-pass characteristics in its transfer function H ( ω ) H(\omega)\, .
For the demonstration of our initial claim, we first take the simplest strategy: the LPF h ( t ) h(t) is a windowed moving average (MA) over a fixed-size window of width w w\, . h ( t ) = λ w Π ( t w − 1 2 ) = { λ / w , if 0 ≤ t < w . 0 , otherwise . h(t)=\frac{\lambda}{w}\Pi\left(\frac{t}{w}-\frac{1}{2}\right)=\begin{cases}\lambda/w\,,&\text{if $0\leq t<w$}\,.\\ 0\,,&\text{otherwise}.\end{cases} (6) The function Π \Pi stands for the Heaviside Pi function, which has value of one in the interval [ − 0.5 , 0.5 ] [-0.5,0.5] and of zero elsewhere. This filter has only two scalar hyperparameters: the scalar factor λ \lambda and the window size w w\, . As shown in Algorithm 1 , we implement this filter h h with a fixed-capacity queue Q Q storing the intermediate parameter updates u ( t ) u(t) into the queue Q Q\, . The average of the queue Q Q is the low-pass filtered gradients which is added to the current parameter update at each optimizer step.
### 2.2 Experiment
We first demonstrate our conjecture on the algorithmic dataset used in the first report on grokking [ Power et al., 2022 ] . The network is a two-layer decoder-only Transformer [ Vaswani et al., 2017 ] trained to predict the answer of a modular binary multiplication operation x ⋅ y ( mod 97 ) x\cdot y\;(\mathrm{mod}\;97)\, . The training curve of this task is shown in Figure 2 as ‘baseline.’ Comparing the time to reach the accuracy of 0.95 0.95\, , the generalization, i.e., the late saturation of the validation accuracy, happens after × 97.3 \times 97.3 iterations after the rapid saturation of the training accuracy (the overfitting). Figure 4 shows empirical proof of effectiveness of Algorithm 1 , our Grokfast-MA algorithm on this task. Choosing the hyperparameters from a simple grid search over λ ∈ { 1 , 2 , 5 , 10 } \lambda\in\{1,2,5,10\} and w ∈ { 2 , 5 , 10 , 20 , 50 , 100 , 200 } w\in\{2,5,10,20,50,100,200\}\, , we found that the filter works best when λ = 5 \lambda=5 and w = 100 w=100\, . As shown in Figure 4 , the number of iterations taken to make the validation accuracy reach 95% of the training accuracy is reduced by × 13.57 \times 13.57\, , which is a remarkable reduction of training iterations.
### 2.3 Discussion
Figure 4 demonstrates high effectiveness of our approach. However, few questions are still left unanswered regarding the modified training dynamics and the combined effect with the weight decay , which is previously shown to be another important algorithmic factor that governs the grokking effect [ Liu et al., 2022b ] . We devise more experiments to answer these questions:
#### Q1. Are both slow and fast gradients necessary?
Our approach is based on our belief that the low-pass filtered gradient updates, or the slow gradients, contribute to the generalization. The most obvious question is then: can we not use the fast gradients and replace the original sequence of gradients with the low-pass filtered components? Using only the slow gradients calculated from a moving average filter in Algorithm 1 is equivalent to using larger, overlapping minibatches. We conduct an experiment with a modified algorithm that replaces the line 8 of Algorithm 1 with g ^ t ← λ ⋅ Avg ( Q ) \hat{g}_{t}\leftarrow\lambda\cdot\mathrm{Avg}(Q)\, . Figure 5 shows the result. 1-stage means that the gradient replacement happens from the beginning of the training, which is set by default, and 2-stage means that the effect of Grokfast happens after the model overfits to the training data at iteration 500 500\, . The results clearly reveal that removing the original gradients leads to much slower and unstable training. In conjunction with the result in Figure 4 , we can conclude that both the fast and the slow components of the gradients are necessary for faster grokking.
#### Q2. Exploiting state transition in the training of a model under grokking.
We can alternatively interpret the training dynamics of a model under the grokking phenomenon as a state transition . In this viewpoint, the model sequentially goes through three distinct stages: ( A \mathrm{A} ) initialized , where both training and validation losses are not saturated, ( B \mathrm{B} ) overfitted , where the training loss is fully saturated but the validation loss is not, and ( C \mathrm{C} ) generalized , where both losses are fully saturated. In the experimental setting of Figure 4 , state transition of A → B \mathrm{A}\rightarrow\mathrm{B} happens roughly after iteration 500 500\, . This interpretation allows us to try out a staged strategy for optimization, where different algorithms are applied to the model during the two transition phases A → B \mathrm{A}\rightarrow\mathrm{B} (from iteration 0 to 499) and B → C \mathrm{B}\rightarrow\mathrm{C} (after iteration 500) as described in Table 1 . Figure 6 and Table 1 summarize the result of the experiment. As the results show, we can accelerate the grokking effect further by × 1.53 \times 1.53 by separating the training stage of the model and applying Grokfast-MA only after the model becomes overfitted, suggesting an adaptive optimizer.
#### Q3. Synergistic effect with weight decay.
Besides from our gradient filtering approach, the authors of Omnigrok [ Liu et al., 2022b ] have suggested that the weight decay hyperparameter is a critical determinant of the grokking phenomenon. According to the report, the grokking phenomenon appears and even becomes faster when the weight decay becomes larger. We, therefore, conduct additional experiments to find out how these two approaches affect the model when applied together. The results are summarized in Figure 7 . Compared with the result from Grokfast-MA with no weight decay ( orange ), applying the weight decay ( red ) generally yields even faster generalization. The maximum acceleration appears at wd = 0.01 \mathrm{wd}=0.01 with × 3.72 \times 3.72 faster generalization than Grokfast-MA with no weight decay. We choose this result of × 50.49 \times 50.49 faster grokking to be our main demonstration in Figure . Interestingly, Figure 7 also reveals that applying the same weight decay with no Grokfast-MA ( brown ) makes the training unstable. The results demonstrates that applying our gradient filtering and setting up a proper weight decay together gives synergistic benefits.
### 2.4 Limitations
Although Algorithm 1 shows highly effective results, it requires w w times more memory to store all the previous gradients, limiting its utilization. Replication of the model parameters also makes the training slower; using w = 100 w=100\, , the training time per iteration is increased by × 2.4 \times 2.4\, measured with a single 1080 Ti GPU. Still, the reduction of wall clock time before the delayed generalization of the results in Figure 7 is × 20.5 \times 20.5\, , which is also a notable reduction of time. Though the computation time does not linearly scale with the memory requirements, Algorithm 1 is not generally applicable to the larger models. This problem leads to our new design utilizing sequential filter in the next section.
## 3 Grokfast with Exponential Moving Average Gradient Filter
In the previous section, we empirically prove that using an LPF to the sequence of model parameter updates leads to faster generalization under the grokking phenomenon. However, for practical purposes, we require an LPF design with a smaller memory footprint. To this end, we modify Algorithm 1 by replacing the windowed moving average with an exponential moving average (EMA) filter. The impulse response of the filter becomes: h ( t ) = λ ( 1 − α ) ∑ τ = 0 t α τ δ ( t − τ ) = λ α t ( 1 − α ) , h(t)=\lambda(1-\alpha)\sum_{\tau=0}^{t}\alpha^{\tau}\delta(t-\tau)=\lambda\alpha^{t}(1-\alpha)\,, (7) where δ ( t ) \delta(t) is the discrete unit impulse at the origin. This filter also has two hyperparameters: the scalar factor λ \lambda and the scalar momentum α \alpha\, . The corresponding Algorithm 2 only requires additional memory with the same size of the model itself, reducing × 50 \times 50 amount of required memory compared to Algorithm 1 . The time and the frequency responses of the filters are compared in Figure 3 .
## 4 Experiment