Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Grokfast: Accelerated Grokking by Amplifying Slow Gradients

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

TL;DR

Grokfast addresses the puzzling grokking phenomenon—delayed generalization long after training accuracy saturates—by treating parameter updates as a stochastic signal and amplifying the slow, generalization-inducing gradient components. It introduces two practical gradient-filter augmentations, Grokfast-MA and Grokfast-EMA, that effectively implement a low-pass frequency boost on gradients (and thus on updates) to hasten the transition from overfitting to generalization. The authors provide theoretical justification showing that gradient filtering before the optimizer preserves the same transfer characteristics as filtering the gradient signal itself, and validate the approach across algorithmic data, MNIST, QM9, and IMDb, achieving large reductions in grokking time and improvements in final metrics. While memory/compute trade-offs exist—especially for the MA variant—the EMA version delivers substantial speedups with a modest footprint and broad compatibility, offering a practical tool for studying and exploiting grokking in diverse settings.

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 $\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.

Grokfast: Accelerated Grokking by Amplifying Slow Gradients

TL;DR

Grokfast addresses the puzzling grokking phenomenon—delayed generalization long after training accuracy saturates—by treating parameter updates as a stochastic signal and amplifying the slow, generalization-inducing gradient components. It introduces two practical gradient-filter augmentations, Grokfast-MA and Grokfast-EMA, that effectively implement a low-pass frequency boost on gradients (and thus on updates) to hasten the transition from overfitting to generalization. The authors provide theoretical justification showing that gradient filtering before the optimizer preserves the same transfer characteristics as filtering the gradient signal itself, and validate the approach across algorithmic data, MNIST, QM9, and IMDb, achieving large reductions in grokking time and improvements in final metrics. While memory/compute trade-offs exist—especially for the MA variant—the EMA version delivers substantial speedups with a modest footprint and broad compatibility, offering a practical tool for studying and exploiting grokking in diverse settings.

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 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.
Paper Structure (31 sections, 3 theorems, 25 equations, 14 figures, 6 tables, 2 algorithms)

This paper contains 31 sections, 3 theorems, 25 equations, 14 figures, 6 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Amplifying the Low-Frequencies of the Stochastic Gradients
  3. Filter Design
  4. Experiment
  5. Discussion
  6. Q1. Are both slow and fast gradients necessary?
  7. Q2. Exploiting state transition in the training of a model under grokking.
  8. Q3. Synergistic effect with weight decay.
  9. Limitations
  10. Grokfast with Exponential Moving Average Gradient Filter
  11. Experiment
  12. Algorithmic Data
  13. MNIST
  14. QM9
  15. IMDb
  16. ...and 16 more sections

Key Result

Theorem A.1

Let $g(t)$ be a scalar signal defined over a discrete time $t \in \{0, 1, \ldots, T\}\,$. Let $h(t)$ be a univariate time-invariant filter defined over the same domain $t\,$. A linear optimizer $O$ is defined as: with scalar coefficients $A, B, C,$ and $D\,$, and $x(0) = g(0)\,$. The output of the system $u(t)$ is, therefore, a function of $g(t)$ and $t\,$, i.e., $u(t) = u(g(t), t)\,$. Let the mo

Figures (14)

  • Figure 1: Grokfast interprets training dynamics of a network parameter as a stochastic signal and amplify the low frequency variations for faster grokking.
  • Figure 2: Accelerating generalization of a model under grokking phenomenon. Our $\textsc{Grokfast}$ is a simple algorithmic modification upon existing optimizers to pull forward the time of event of sudden generalization after overfitting, also known as the grokking phenomenon.
  • Figure 3: Time and frequency domain plots of the gradient filters. Figures (a, b) and (c, d) depict the impulse responses and the transfer functions of the filters of Algorithm \ref{['alg:avg']} and \ref{['alg:main']}, i.e., the MA and the EMA filters $h(t)\,$. We treat training iterations as discrete timesteps.
  • Figure 4: Acceleration of delayed generation with Grokfast-MA. The amount of acceleration relies on the two hyperparameters, amplifier gain $\lambda$ and window size $w\,$. Each hyperparameter has a sweet spot; increasing one arbitrarily does not guarantee faster acceleration. Figures \ref{['fig:ma:acc_gain']} and \ref{['fig:ma:loss_gain']} use $w = 100$ except the baseline. Figures \ref{['fig:ma:acc_window']} and \ref{['fig:ma:loss_window']} use $\lambda = 5$ except the baseline.
  • Figure 5: Although adding the slow component of the gradients is effective in accelerating grokking, the slow component cannot be used alone as a replacement.
  • ...and 9 more figures

Theorems & Definitions (6)

  • Theorem A.1
  • proof : Proof of Theorem \ref{['thm:linear_optim']}
  • Proposition A.2: SGD with momentum
  • proof : Proof of Proposition \ref{['prop:sgdm']}
  • Proposition A.3: SGD with Nesterov's momentum
  • proof : Proof of Proposition \ref{['prop:nesterov']}