PaperlayerView exact revision on arXiv

Source-linked AI summary

Grokfast: Accelerated Grokking by Amplifying Slow Gradients

Jaerin Lee, Bong Gyun Kang, Kihoon Kim, Kyoung Mu Lee

arXiv:2405.20233v2cs.LGcs.AI

TL;DR

Grokking is delayed generalization after a model overfits its training data, and the delay carries a computational burden. This paper treats parameter updates as signals and amplifies low-frequency gradients; experiments report up to ×50 faster grokking across diverse tasks. Its moving-average filter has a memory cost that limits use on larger models, motivating an exponential moving-average design.

  • Problem

    Grokking delays generalization until long after training-data overfitting, creating a computational burden for practitioners.

    Source: §1 · p. 1

  • Method

    The method amplifies low-frequency gradient components, using a low-pass filter added to the original gradient signal.

    Source: §1 · p. 2, §2.1 · p. 3

  • Results

    Experiments across varied architectures and algorithmic, image, language, and graph tasks report up to ×50 faster grokking.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    The authors conclude that state-space and frequency representations of training dynamics are useful for studying grokking.

    Source: §7 · p. 10

    The moving-average filter requires w times more memory and is not generally applicable to larger models, motivating the sequential-filter design.

    Source: §2.4 · pp. 5–6

Abstract

from arXiv · show

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 $\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 delays generalization until long after training data are overfit, while its causes and sufficient conditions remain incompletely characterized. This work interprets parameter updates as signals and proposes amplifying low-frequency gradients to accelerate that delayed generalization.

1 Introduction

  • Grokking is delayed generalization after training-data overfitting, and its causes and sufficient conditions remain uncharacterized.Prior theory-oriented work has related it to double descent [Davies et al., 2023, Huang et al., 2024].Source: §1 · p. 1
  • Prior reports describe generalization occurring more than tenfold training iterations after overfitting, a computational burden that limits practical appeal.The paper takes a practitioner’s standpoint and aims to accelerate grokking.Source: §1 · p. 1
  • In an example grokking curve, validation-loss dynamics are a few orders of magnitude slower than training-loss dynamics, suggesting fast updates contribute to overfitting and slow updates to generalization.The parameter changes are treated as a direct consequence of the changes in the losses.Source: §1 · p. 2
  • The approach treats each parameter’s optimizer update over training iterations as a discrete random signal, then analyzes its frequency-domain representation.Training iterations serve as the time index; the update function abstracts iterative optimizers including SGD variants.Source: §1 · p. 2
  • The authors hypothesize that amplifying low-frequency components of the gradient signal accelerates generalization under grokking.They relate the hypothesis to slow-varying parameter updates and, for first-order optimizers, the linear correlation between gradients and updates.Source: §1 · p. 2
  • The method is reported to produce up to ×50 faster grokking across experiments involving Transformers, MLPs, RNNs, (Graph-)ConvNets, and algorithmic, image, language, and graph tasks.It is described as requiring only a few lines of additional code and as applicable to most machine-learning frameworks.Source: §1 · p. 2

2 Amplifying the Low-Frequencies of the Stochastic Gradients

GROKFAST-MA adds a low-pass-filtered gradient component to the original update to accelerate grokking. Experiments show faster generalization, while retaining the original gradients and managing memory and hyperparameters remain important.

2.1 Filter Design

  • The moving-average filter adds a low-pass-filtered gradient signal to the current update, with gain λ and window width w as its two scalar hyperparameters.The filter averages updates over a fixed-size window; each hyperparameter has a sweet spot, so increasing either does not guarantee more acceleration.Source: §2.1 · p. 3, Figure 4 · p. 4

2.2 Experiment

  • ×13.57 fewer iterations were needed to reach 95% validation accuracy relative to training accuracy with GROKFAST-MA on modular multiplication, using λ = 5 and w = 100.The baseline reached that validation threshold ×97.3 iterations after training accuracy rapidly saturated.Source: §2.2 · p. 4, Figure 4 · p. 4

2.3 Discussion

  • Replacing the original gradients with only the slow, moving-average component leads to slower and less stable training, so both fast and slow components are necessary for faster grokking.The slow-only variant is equivalent to using larger, overlapping minibatches.Source: Figure 5 · p. 4, §2.3 · p. 5
  • Applying GROKFAST-MA only after overfitting, at roughly iteration 500, further accelerates grokking by ×1.53 compared with the default strategy.The authors frame this staged schedule as suggesting an adaptive optimizer.Source: §2.3 · p. 5, Figure 6 · p. 5
  • Appropriate weight decay greatly enhances GROKFAST-MA's acceleration, but weight decay alone does not always help.Source: Figure 7 · p. 5

2.4 Limitations

  • Algorithm 1 requires w times more memory to store prior gradients, and with w = 100 its iteration time increased ×2.4 on a single 1080 Ti GPU.The authors say the method is not generally applicable to larger models, motivating a sequential-filter design.Source: §2.4 · pp. 5–6
  • Despite that iteration-time increase, the reported wall-clock time before delayed generalization fell ×20.5 in the Figure 7 results.Source: §2.4 · p. 6

3 Grokfast with Exponential Moving Average Gradient Filter

GROKFAST replaces the windowed moving average with an exponential moving average (EMA) filter to reduce the filter’s memory footprint.

3 Grokfast with Exponential Moving Average Gradient Filter

  • The EMA filter replaces the windowed moving average to provide a lower-memory low-pass filter.The filter’s impulse response is given for this EMA design.Source: §3 · p. 6
  • The filter’s impulse response defines the EMA design used in the modified algorithm.The passage introduces this response after describing the replacement of the windowed moving average.Source: §3 · p. 6
  • Algorithm 2 uses only model-sized additional memory, reducing required memory by 50× compared with Algorithm 1.The EMA filter has two hyperparameters: scalar factor λ and scalar momentum α.Source: §3 · p. 6

4 Experiment

GROKFAST is evaluated on algorithmic, image, graph, and language tasks that exhibit grokking. Across the supplied experiments, it accelerates delayed generalization or convergence, while the effects and suitable settings vary by task.

4 Experiment

  • The experiments extend grokking acceleration beyond algorithmic data to conventional tasks with larger and more complex datasets.The evaluated task types include images, graphs, and language.Source: §4 · pp. 6–7

4.1 Algorithmic Data

  • On modular multiplication, the EMA filter captures slow gradient variation with 50× less computational resource use than Algorithm 1, and GROKFAST shows a smaller validation-loss peak.The authors interpret the smaller peak as keeping the generalization route B→C close to the global optimum at state C.Source: §4.1 · p. 7
  • A cutoff-frequency sweet spot appears across the reported experiments; the authors recommend λ ∈ [0.1, 5] and α ∈ [0.8, 0.99], while weight decay depends on the task.For the EMA filter, α corresponds to the cutoff frequency alongside the window size w for the moving-average filter.Source: §4.1 · p. 7

4.2 MNIST

  • The MNIST delay until grokking falls 22.0×, and final evaluation accuracy rises from 89.8% to 91.2%.The experiment uses a three-layer ReLU-MLP with α = 0.8, λ = 0.1, and wd = 2.0.Source: §4.2 · p. 7, Figure 9 · p. 7

4.4 IMDb

  • On IMDb, GROKFAST produces faster generalization with higher validation error and lower validation loss.The experiment uses a two-layer LSTM; the curves shown are exponential moving averages with momentum 0.95.Source: Figure 11 · p. 8, §4.4 · p. 8

4.2 MNIST

  • On QM9, GROKFAST makes validation loss drop faster and by a larger margin than the baseline.The comparison uses a GCNN and α = 0.9, λ = 1.0, and wd = 0.01; this task has no accuracy measure for comparing convergence speed.Source: Figure 10 · p. 7, §4.2 · p. 8

5 More Discussion

GROKFAST adds a smoothed-gradient residual before the optimizer, unlike typical momentum updates. Its parameter-space trajectories differ from the baseline, supporting an interpretation of grokking as transitions among distinct states.

5.1 Difference between Algorithm 2 and the Momentum in Typical Optimizers

  • GROKFAST adds the smoothed gradient as a residual before the optimizer, rather than using scaled momentum as the parameter update.The filtering is applied independently of the underlying first-order gradient-based optimizer.Source: §5.1 · p. 8

5.2 Visualizing Trajectories

  • The baseline’s path from overfitting to generalization is longer than its path from initialization to overfitting; GROKFAST makes the AB and BC distances more even and shortens AC by 16×.The authors interpret the states as distinct, even though final-state accuracy and losses are similar between methods.Source: §5.2 · p. 9, Table 2 · p. 9
  • At overfitting, GROKFAST deviates 8× farther from initialization than the baseline, with 5× smaller standard deviation in distances from initialization.The authors suggest the two overfitted states may comprise distinct parameter sets with possibly different topologies.Source: §5.2 · p. 9, Table 2 · p. 9
  • GROKFAST’s distance variances are reported as hundredfold smaller, but the source and role of this determinism and its benefits remain unresolved.Using only slow gradients while neglecting the original gradients leads to instability.Source: §5.2 · p. 9

6 Related Work

Prior work relates grokking to optimization, regularization, model and dataset properties, and internal learning progress. GROKFAST builds on this literature with an optimizer augmentation and a state-transition interpretation.

6 Related Work

  • Prior accounts link grokking to double descent, pattern-learning speed, model and dataset sizes, and internal progress that may not be abrupt.The cited discussion also connects grokking to optimization and weight decay.Source: §6 · pp. 9–10
  • Earlier studies identify optimizer choice, batch training, weight decay, noise, dropout, and learning rate as factors affecting grokking patterns.Prior findings differ on whether explicit regularization is necessary; one account attributes grokking without it to an optimizer’s implicit regularization.Source: §6 · p. 10
  • GROKFAST is presented as an augmentation for existing first-order optimizers, motivated by a state-space model and dual-domain view of training dynamics.The authors frame it as a tool for practical study of grokking.Source: §6 · p. 10

7 Conclusion

The paper interprets parameter updates as signals with fast- and slow-varying components and reports that amplifying the slow component can bring forward grokking. Its appendix shows that, for linear optimizers, filtering gradients is equivalent to applying the same filter to parameter updates.

7 Conclusion

  • Amplifying slow-varying gradient components with low-pass filtering can bring forward sudden late generalization, reducing required training iterations by up to ×50.The authors also argue that the frequency representation of training dynamics is useful for studying grokking.Source: §7 · p. 10

A Frequency Responses of the Parameter Updates under Grokfast

  • The appendix models optimizer updates as outputs of a state-space system driven by gradients, with the state recursively updated from the previous state and current gradient.The update output depends on the state and gradient; the appendix applies this setup to SGD with momentum and Nesterov momentum.Source: A Frequency Responses of the Parameter… · p. 12
  • For a linear optimizer and a scalar, time-invariant filter, filtering gradients is equivalent to applying the same filter to parameter updates.This result supports using gradient filtering to produce low-pass parameter updates.Source: A Frequency Responses of the Parameter… · pp. 12–13

B Task Details

The appendix summarizes implementation details for the tasks in Section 4 and points readers to the PyTorch implementation.

B Task Details

  • The appendix summarizes implementation details for each task addressed in Section 4.Readers can also consult the official PyTorch implementation.Source: B Task Details · p. 14
  • The authors direct readers to their official implementation in PyTorch.The cited framework is PyTorch [Paszke et al., 2019].Source: B Task Details · p. 14
  • The implementation reference complements the task descriptions with code readers can consult.The passage identifies the implementation as official.Source: B Task Details · p. 14

B.1 Binary Operation (Algorithmic Data)

The binary-operation experiments use modular multiplication with p = 97 and a two-layer decoder-only Transformer trained with cross-entropy and Adam.

B.1 Binary Operation (Algorithmic Data)

  • The algorithmic task predicts modular multiplication x · y (mod 97) using a two-layer decoder-only Transformer.The model has hidden dimension 128, four attention heads, and positional embedding length five.Source: §B.1 · p. 14
  • Training uses cross-entropy loss and Adam with betas (β1, β2) = (0.9, 0.98), learning rate 10^-3, and batch size 512.The learning rate is warmed up linearly over the first 10 iterations.Source: §B.1 · p. 14
  • The Transformer uses GELU and layer normalization, followed by layer normalization and a linear output layer to produce logits.The network is based on the modular binary-operation setup first reported for grokking [Power et al., 2022].Source: §B.1 · p. 14

B.2 MNIST

The task details cover MNIST classification, QM9 polarizability estimation, and IMDb sentiment classification, each configured to simulate or study grokking under a specified training setup.

B.2 MNIST

  • MNIST uses a three-layer, width-200 ReLU MLP trained on 1,000 images with AdamW for up to 10^5 iterations.The model uses ×8 Kaiming weight initialization, flattened 28 × 28 images, one-hot labels, and MSE loss.Source: §B.2 · pp. 14–15

B.3 QM9

  • QM9 experiments estimate isotropic polarizability from the first 100 molecules using a Graph ConvNet and a 50:50 train-validation split.The model uses AdamW, batch size 32, ×3 Kaiming weight initialization, and 50k training iterations.Source: §B.3 · p. 15

B.4 IMDb

  • IMDb sentiment classification uses a two-layer LSTM on the first 1,000 reviews, split 75:25, with training stopped at 10k iterations.Reviews are tokenized to the 1,000 most frequent words and padded to length 500; training uses binary cross-entropy and AdamW.Source: §B.4 · p. 15

C Autograd Implementation

The implementation adds gradient-filtering helpers for Grokfast and Grokfast-MA to an autograd training loop. The authors state that only two code lines need modification to use the helper.

C Autograd Implementation

  • Grokfast uses an exponential moving average of gradients, while Grokfast-MA stores gradients in a bounded history for filtering.The EMA helper combines the stored gradient with the current gradient using alpha; the MA helper appends detached gradients to a deque.Source: C Autograd Implementation · pp. 15–16
  • The training loop calls either Grokfast or Grokfast-MA after backpropagation and before the optimizer step.The two options expose different arguments: alpha and lamb for Grokfast, or window_size and lamb for Grokfast-MA.Source: C Autograd Implementation · p. 16
  • Only two lines of the training code are modified to apply the helper between gradient calculation and the optimizer call.The authors describe the implementation using PyTorch autograd.Source: C Autograd Implementation · pp. 15–16

D Time and Memory Requirements

The section reports that Grokfast’s additional computational burden is offset by accelerated delayed generalization, with negligible extra VRAM use. Further visualizations show Grokfast trajectories reaching more consistent generalization states closer to initialization.

D Time and Memory Requirements

  • The additional computational burden is described as compensated by accelerated delayed generalization, while VRAM requirements are negligible compared with baseline.The reported time and memory measurements used a single GTX 1080 Ti GPU.Source: D Time and Memory Requirements · pp. 16–17

E More Visualization

  • The visualization projects normalized parameter differences from checkpoints onto the two most salient directions of variation and marks initialization, overfitting, and generalization states.The plotted states include overfitting at 500 iterations.Source: E More Visualization · pp. 17–18
  • Grokfast trajectories reach generalization states closer to initialization, with less variation across runs than the baseline trajectories.The baseline runs’ generalization states vary substantially, while Grokfast trajectories are more consistent.Source: Figure 13 · p. 17, E More Visualization · p. 18

Related papers

Loading 2405.20233v2…