Grokfast: Gradient filters for faster grokking

• This paper will make a much better contribution if the authors provide a rigorous mathematical derivation or analysis for presenting intuition about why they believe amplifying the low-frequency component of $G(\omega)$ accelerates the speed of generalization under grokking phenomenon.
• It’s not clear enough about how scaling slow gradients will benefit or hurt the grokking process. Can the occurrence of the grok phenomenon be slowed down by reducing the low-frequency component of updates? If possible, I recommend the authors add this experiment to verify their hypothesis by using some techniques like high-pass filtering the gradients.
• Although Grokfast performs fairly better than baselines in the experiment, I doubt its effectiveness in training the modern large-scale neural networks since the MA and EMA tricks are now widely used in optimizers, but none of them could beat popular SGDm and AdamW. We don’t know when and whether the grok phenomenon will appear. For example, grokking might never appear in pretraining language model on comprehensive dataset. I am concerned that the conclusion and method might not be able to generalize to all ML tasks.
• Minor issue: The authors mixed up the labels of algorithms in Appendix C.

In most setups, grokking can be avoided by selecting appropriate hyperparameters. Liu et al. (2022a) describe grokking as an interesting but rare phenomenon, noting that simply optimizing hyperparameters, such as the learning rate, can significantly speed up training and prevent grokking without requiring special algorithms. While the authors of this paper do acknowledge the importance of optimizer choice and hyperparameters in principle, they emphasize their algorithm’s ability to accelerate generalization with minimal code changes (e.g., lines 19-20, 196-197, 473, 518, 526) and present speedup as their primary goal (lines 14-15). However, given the existing literature, achieving faster training or avoiding grokking through these methods is not particularly surprising. In general, itmight be an interesting question how different hyperparameters and optimizers affect grokking, but the authors do not present their results in this way.

As the authors note, their exponential moving average (EMA) algorithm closely resembles a conventional optimizer that incorporates a momentum hyperparameter. Although the authors also introduce an algorithm utilizing a simple moving average (MA), which demonstrates a slight improvement in generalization speed up (50x for MA compared to 44x for EMA), the overall impact on training performance appears comparable to that of the exponential moving average (see Figure 2 and Figure 4, panels (a) and (b)). This observation raises the question, why no conventional momentum hyperparameter was investigated and compared to, and leaves the reader with results solely for the unconventional algorithms and no proper comparison.

Overall, the authors provide minimal hyperparameter comparison for the baseline case that does not incorporate any of their proposed algorithms. In particular, they do not investigate the effects of varying the learning rate, which has been shown to significantly influence the grokking phenomenon (Liu et al., 2022a).

Liu et al. (2022a): https://arxiv.org/abs/2205.10343

Minor comment:

line 188: Typo: mentioning window size $w$ as a hyperparameter of the exponential moving average algorithms instead of momentum $\alpha$.

1. The approach developed in the paper is motivated by the claim that delayed generalization in grokking is due to high-frequency oscillations of the gradient (line 94). I don't see what confirms this claim. The paper does not seem to cite or present any evidence here.

2. The eficiency of the proposed filtering approach is demonstrated only by experimental evidence. The paper does not make any theoretical analysis beyond the vague generic statement that low-pass filtering should reduce high-frequence oscillations and thus reduce generalization delay.

3. The experiments with synthetic problems (modular multiplication) do confirm a substantial reduction of generalization delay by the proposed method. However, in the case of experiments with real data the situation is not so obvious. For MNIST, the authors seem to use an oddly inefficient three-layer ReLU-MLP that has a very low accuracy of 90%. Their algorithms speeds-up generalization, but the accuracy remains at around 91%. In contrast, the usual accuracy of even basic MLP's on MNIST is 98% or so. I think that it is misleading to omit a discussion of this point in the paper.
   In the case of QM9 and IMDB both grokking and the effect of the proposed algorithm are not very pronounced. The proposed filtering seems to be analogous to adding momentum or averaging - commonly used techniques well-known to benefit optimization. The improvement visible in the plots can be attributed to standard improvement from these techniques.

4. The paper contains a theoretical Appendix A, but it only includes theorems that seem to be of relatively technical nature (equivalence between filtering the gradients and the parameter updates) and don't do much to increase confidence in the proposed method. These theorems don't seem to be referenced in the main text.

5. Given its main focus, the paper does not provide sufficient background on phenomenology and previous theoretical studies of grokking. The paper includes Section 6 on related work, but it only relatively superficially mentions previous research, in particular not going into much detail how various factors increase or reduce grokking. I would expect the present paper to discuss much more extensively how its idea of low-pass filtering corresponds to the existing phenomenology (ideally, already in the introduction) and alternative theories of grokking. In particular, I have also seen the paper [1] explaining grokking by the lack of feature learning and claiming an alternative approach to remove grokking by inducing feature learning.

[1] Kumar et al., Grokking as the Transition from Lazy to Rich Training Dynamics, ICLR 2024

6. Grokking appears to be a relatively rare and somewhat artificial phenomenon, uncommon in real-world data and neural networks. Accordingly, whereas a better theoretical understanding of grokking would certainly be appreciated, I'm not sure that algorithms aimed at eliminating grokking are practically very important.

Summarizing, I think that the paper in its present state is not good enough for ICLR in terms of significance, quality, and usefulness for the reader. I think that the paper can be substantially improved by addressing the items indicated above.

• The moving average in Algorithm 1 to compute hat g_t is conceptually very close to the Adam method. In the article, the authors have used this method (AdamW) in the numerical results, therefore it is unclear whether the effect of moving averaging is so crucial or not to explain the delay-time reduction.

In this sense, the main message of the article is conceptually not fully convincing. I suggest the authors to conduct an ablation study on the same experiments without using Adam optimizers (e.g. u(g,t) is simply SGD without momentum).

• The writing could be further improved. Certain notation such as the convolution in eq. 3 is not standard. One would use h*g(t) in this case.