Grokfast: Accelerated Grokking by Amplifying Slow Gradients — part 3/4
Cleaned-text code points [31500, 47384). Bibliography omitted here; source.md includes it.
To compare the difference between the frequency responses of the parameter update u ( t ) u(t) and of the modified update u ^ ( t ) \hat{u}(t) in equation ( 4 ), we can think of an equivalent filter h ^ ( t ) \hat{h}(t) defined to satisfy the following relationship in addition to equations ( 3 ) and ( 4 ): u ^ ( t ) = u ( g ^ ( t ) , t ) = u ( g ( t ) + h ( t ) ∗ g ( t ) , t ) = u ( t ) + h ^ ( t ) ∗ u ( t ) . \hat{u}(t)=u(\hat{g}(t),t)=u(g(t)+h(t)*g(t),t)=u(t)+\hat{h}(t)*u(t)\,. (10) From our assumption of the linear time-invariant, scalar filters h ( t ) h(t) and the linear optimizer, we can deduce the equivalence between h ( t ) h(t) and h ^ ( t ) \hat{h}(t)\, . The following theorem is a generalized claim that applies to any SGD-based first-order optimizers including Nesterov’s momentum.
###### Theorem A.1 .
Let g ( t ) g(t) be a scalar signal defined over a discrete time t ∈ { 0 , 1 , … , T } t\in\{0,1,\ldots,T\}\, . Let h ( t ) h(t) be a univariate time-invariant filter defined over the same domain t t\, . A linear optimizer O O is defined as: x ( t ) \displaystyle x(t) = A x ( t − 1 ) + B g ( t ) , t > 0 , \displaystyle=Ax(t-1)+Bg(t)\,,\quad t>0\,, (11) u ( t ) \displaystyle u(t) = C x ( t ) + D g ( t ) , t ≥ 0 , \displaystyle=Cx(t)+Dg(t)\,,\qquad\,\,\,\,t\geq 0\,, (12) with scalar coefficients A , B , C , A,B,C, and D D\, , and x ( 0 ) = g ( 0 ) x(0)=g(0)\, . The output of the system u ( t ) u(t) is, therefore, a function of g ( t ) g(t) and t t\, , i.e., u ( t ) = u ( g ( t ) , t ) u(t)=u(g(t),t)\, . Let the modified input g ^ ( t ) \hat{g}(t)\, , the modified output u ^ ( t ) \hat{u}(t)\, , and the equivalent filter h ^ ( t ) \hat{h}(t) be defined to satisfy the equations ( 3 ) and ( 10 ). Then, h ^ ( t ) = h ( t ) , \hat{h}(t)=h(t)\,, (13) for t ∈ { 0 , 1 , … , T } t\in\{0,1,\ldots,T\}\, .
###### Proof of Theorem A.1 .
For simplicity, we first adopt discrete-time Fourier transform over t ∈ ℤ t\in\mathbb{Z}\, . That is, we assume that the signals are defined across every positive and negative integer t t\, . Since the value of u ( t ) u(t) can be defined arbitrarily outside the interval [ 0 , T ] [0,T] without modifying the optimization algorithm, we can manually assign g ( t ) g(t) and x ( t ) x(t) for t ∉ [ 0 , T ] t\notin[0,T] as: g ( t ) \displaystyle g(t) = 0 t ∉ [ 0 , T ] , \displaystyle=0\qquad\qquad\qquad\quad\,\,\,t\notin[0,T]\,, (14) x ( t ) \displaystyle x(t) = { A t ( 1 − B ) g ( 0 ) t < 0 , A t − T x ( T ) t > T . \displaystyle=\begin{cases}A^{t}(1-B)g(0)&t<0\,,\\ A^{t-T}x(T)&t>T\,.\\ \end{cases} (15) Then, equations ( 11 ) and ( 12 ) hold for t ∉ [ 0 , T ] t\notin[0,T]\, . Note that to make the optimizer O O stable, the scalar coefficient A A should satisfy 0 < A < 1 0<A<1\, . Therefore, the signals g ( t ) g(t) and x ( t ) x(t) are well-defined.
Consider a discrete-time Fourier transform ℱ \mathcal{F} defined as: ℱ { f ( t ) } ( ω ) = ∑ t = − ∞ ∞ f ( t ) e − i ω t . \mathcal{F}\{f(t)\}(\omega)=\sum_{t=-\infty}^{\infty}f(t)e^{-i\omega t}\,. (16) In the frequency domain, with G ( ω ) = ℱ { g ( t ) } G(\omega)=\mathcal{F}\{g(t)\}\, , U ( ω ) = ℱ { u ( t ) } U(\omega)=\mathcal{F}\{u(t)\}\, , and X ( ω ) = ℱ { x ( t ) } X(\omega)=\mathcal{F}\{x(t)\}\, , the optimizer O O can be equivalently represented as: X ( ω ) \displaystyle X(\omega) = A e − i ω X ( ω ) + B G ( ω ) , \displaystyle=Ae^{-i\omega}X(\omega)+BG(\omega)\,, (17) U ( ω ) \displaystyle U(\omega) = C X ( ω ) + D G ( ω ) . \displaystyle=CX(\omega)+DG(\omega)\,. (18) We can obtain the transfer functions H in-state H_{\text{in-state}} and H in-out H_{\text{in-out}} that converts G G to X X and then to U U\, : H in-state ( ω ) \displaystyle H_{\text{in-state}}(\omega) ≔ X ( ω ) G ( ω ) = B 1 − A e − i ω , \displaystyle\coloneqq\frac{X(\omega)}{G(\omega)}=\frac{B}{1-Ae^{-i\omega}}\,, (19) H in-out ( ω ) \displaystyle H_{\text{in-out}}(\omega) ≔ U ( ω ) G ( ω ) = C X ( ω ) G ( ω ) + D = B C 1 − A e − i ω + D . \displaystyle\coloneqq\frac{U(\omega)}{G(\omega)}=C\frac{X(\omega)}{G(\omega)}+D=\frac{BC}{1-Ae^{-i\omega}}+D\,. (20) If the input g ( t ) g(t) is filtered with a convolutional filter h ( t ) h(t) and then added to itself as equations ( 3 ) and ( 5 ), the state x ( t ) x(t) and the output u ( t ) u(t) of the optimizer O O is changed accordingly while keeping equations ( 11 ) and ( 12 ) hold. We denote x ^ ( t ) \hat{x}(t) and u ^ ( t ) \hat{u}(t) as the modified state and output of the system and X ^ ( ω ) \hat{X}(\omega) and U ^ ( ω ) \hat{U}(\omega) as their spectra. If the filter h ( t ) h(t) is causal, that is h ( t ) = 0 h(t)=0 for t < 0 t<0\, , then we can similarly let x ^ ( 0 ) = g ^ ( 0 ) \hat{x}(0)=\hat{g}(0) and replace x ( t ) {x}(t) and u ( t ) {u}(t) with x ^ ( t ) \hat{x}(t) and u ^ ( t ) \hat{u}(t) in equations ( 14 ) and ( 15 ) to define an IIR system suitable for the infinite-window discrete-time Fourier transform ℱ \mathcal{F}\, : X ^ ( ω ) \displaystyle\hat{X}(\omega) = A e − i ω X ^ ( ω ) + B G ^ ( ω ) , \displaystyle=Ae^{-i\omega}\hat{X}(\omega)+B\hat{G}(\omega)\,, (21) U ^ ( ω ) \displaystyle\hat{U}(\omega) = C X ^ ( ω ) + D G ^ ( ω ) . \displaystyle=C\hat{X}(\omega)+D\hat{G}(\omega)\,. (22) Since the coefficients of the linear systems are the same, the transfer functions are identical: H ^ in-state ( ω ) \displaystyle\hat{H}_{\text{in-state}}(\omega) ≔ X ^ ( ω ) G ^ ( ω ) = B 1 − A e − i ω ≡ H in-state ( ω ) , \displaystyle\coloneqq\frac{\hat{X}(\omega)}{\hat{G}(\omega)}=\frac{B}{1-Ae^{-i\omega}}\qquad\equiv H_{\text{in-state}}(\omega)\,, (23) H ^ in-out ( ω ) \displaystyle\hat{H}_{\text{in-out}}(\omega) ≔ U ^ ( ω ) G ^ ( ω ) = B C 1 − A e − i ω + D ≡ H in-out ( ω ) . \displaystyle\coloneqq\frac{\hat{U}(\omega)}{\hat{G}(\omega)}=\frac{BC}{1-Ae^{-i\omega}}+D\equiv H_{\text{in-out}}(\omega)\,. (24)
From equation ( 5 ), the transfer function of the filter H amp ( ω ) H_{\text{amp}}(\omega) is: H amp ( ω ) = G ^ ( ω ) G ( ω ) = 1 + H ( ω ) , H_{\text{amp}}(\omega)=\frac{\hat{G}(\omega)}{G(\omega)}=1+H(\omega)\,, (25) where H ( ω ) = ℱ { h ( t ) } H(\omega)=\mathcal{F}\{h(t)\}\, . The equivalent post-filter h ^ ( t ) \hat{h}(t) defined by equation ( 10 ) gives another transfer function between the outputs u ( t ) u(t) and u ^ ( t ) \hat{u}(t) of the system: H ^ amp ( ω ) = U ^ ( ω ) U ( ω ) = 1 + H ^ ( ω ) . \hat{H}_{\text{amp}}(\omega)=\frac{\hat{U}(\omega)}{U(\omega)}=1+\hat{H}(\omega)\,. (26) From equations ( 20 ) and ( 24 ), we have: H ^ amp ( ω ) = U ^ ( ω ) U ( ω ) = G ^ ( ω ) G ( ω ) = H amp ( ω ) . \hat{H}_{\text{amp}}(\omega)=\frac{\hat{U}(\omega)}{U(\omega)}=\frac{\hat{G}(\omega)}{G(\omega)}=H_{\text{amp}}(\omega)\,. (27) Therefore, we get: H ^ ( ω ) ≡ H ( ω ) . \hat{H}(\omega)\equiv{H}(\omega)\,. (28) This completes the proof. ∎
In other words, applying any filter h ( t ) h(t) to the sequence of gradients g ( t ) g(t) is equivalent to the same filter h ( t ) h(t) applied to the parameter update u ( t ) u(t) for any linear optimizer O O\, . This implies that a low-pass gradient filter h ( t ) h(t) guarantees the same low-pass property in the modified parameter update signal u ^ ( t ) \hat{u}(t)\, . In many off-the-shelf autograd packages such as PyTorch [ Paszke et al., 2019 ] , filtering the gradients is easier and more straightforward than filtering the intermediate parameter updates. The former only adds a few more lines to the outermost application code 1 1 1 See our implementation at https://github.com/ironjr/grokfast . , whereas the latter requires full implementation of the dedicated optimizer object. Note that the above proof holds regardless of the design of the filter h ( t ) h(t) unless there exists a one-to-one correspondence between H amp H_{\text{amp}} and H H\, .
The followings are direct consequences of the Theorem A.1 .
###### Proposition A.2 (SGD with momentum) .
Let t ∈ { 0 , 1 , … , T } t\in\{0,1,\ldots,T\} be a discrete timestep. Let g ( t ) g(t) be a sequence of gradients of a parameter θ \theta sampled from a stochastic machine learning framework g ( t ) ∼ M ( θ ( t ) , t ) g(t)\sim M(\theta(t),t) and a stochastic gradient descent optimizer O ( μ , τ , η ) O(\mu,\tau,\eta) with a parameter update function u ( g ( t ) , t ) = u ( t ) = θ ( t + 1 ) − θ ( t ) u(g(t),t)=u(t)=\theta(t+1)-\theta(t)\, , a momentum μ \mu\, , a damping constant τ \tau\, , and a learning rate η \eta\, . The parameter update u ( t ) u(t) is, therefore, defined as: m ( t ) \displaystyle m(t) = μ m ( t − 1 ) + ( 1 − τ ) g ( t ) , \displaystyle=\mu m(t-1)+(1-\tau)g(t)\,, (29) u ( t ) \displaystyle u(t) = − η m ( t ) , \displaystyle=-\eta m(t)\,, (30) with a scalar momentum term m ( t ) m(t) for each parameter θ \theta with m ( 0 ) = g ( 0 ) m(0)=g(0)\, . Let h ( t ) h(t) be a scalar, time-invariant, convolutional gradient filter. Let the modified input g ^ ( t ) \hat{g}(t)\, , the modified output u ^ ( t ) \hat{u}(t)\, , and the equivalent filter h ^ ( t ) \hat{h}(t) be defined to satisfy the equations ( 3 ) and ( 10 ). Then, h ^ ( t ) = h ( t ) , \hat{h}(t)=h(t)\,, (31) for t ∈ { 0 , 1 , … , T } t\in\{0,1,\ldots,T\}\, .
###### Proof of Proposition A.2 .
Let A = μ A=\mu\, , B = 1 − τ B=1-\tau\, , C = − η C=-\eta\, , and D = 0 D=0\, . By Theorem A.1 , equation ( 31 ) holds. ∎
###### Proposition A.3 (SGD with Nesterov’s momentum) .
Let t ∈ { 0 , 1 , … , T } t\in\{0,1,\ldots,T\} be a discrete timestep. Let g ( t ) g(t) be a sequence of gradients of a parameter θ \theta sampled from a stochastic machine learning framework g ( t ) ∼ M ( θ ( t ) , t ) g(t)\sim M(\theta(t),t) and a stochastic gradient descent optimizer O ( μ , τ , η ) O(\mu,\tau,\eta) with a parameter update function u ( g ( t ) , t ) = u ( t ) = θ ( t + 1 ) − θ ( t ) u(g(t),t)=u(t)=\theta(t+1)-\theta(t)\, , a momentum μ \mu\, , a damping constant τ \tau\, , and a learning rate η \eta\, . The parameter update u ( t ) u(t) is, therefore, defined as: m ( t ) \displaystyle m(t) = μ m ( t − 1 ) + ( 1 − τ ) g ( t ) , \displaystyle=\mu m(t-1)+(1-\tau)g(t)\,, (32) u ( t ) \displaystyle u(t) = − η ( g ( t ) + μ m ( t ) ) , \displaystyle=-\eta(g(t)+\mu m(t))\,, (33) with a scalar momentum term m ( t ) m(t) for each parameter θ \theta with m ( 0 ) = g ( 0 ) m(0)=g(0)\, . Let h ( t ) h(t) be a scalar, time-invariant, convolutional gradient filter. Let the modified input g ^ ( t ) \hat{g}(t)\, , the modified output u ^ ( t ) \hat{u}(t)\, , and the equivalent filter h ^ ( t ) \hat{h}(t) be defined to satisfy the equations ( 3 ) and ( 10 ). Then, h ^ ( t ) = h ( t ) , \hat{h}(t)=h(t)\,, (34) for t ∈ { 0 , 1 , … , T } t\in\{0,1,\ldots,T\}\, .
###### Proof of Proposition A.3 .
Let A = μ A=\mu\, , B = 1 − τ B=1-\tau\, , C = − η μ C=-\eta\mu\, , and D = − η D=-\eta\, . By Theorem A.1 , equation ( 34 ) holds. ∎
## Appendix B Task Details
For completeness, this section summarizes the implementation details of each task dealt in Section 4 . The readers can also consult our official implementation in PyTorch [ Paszke et al., 2019 ] .
### B.1 Binary Operation (Algorithmic Data)
This is the description of algorithmic data used throughout the manuscript. Following the first report on the grokking phenomenon [ Power et al., 2022 ] , we demonstrate our acceleration algorithms with a binary operation x ⋅ y ( mod p ) x\cdot y\;(\mathrm{mod}\;p)\, , with p = 97 p=97\, . The network is a two-layer decoder-only Transformer [ Vaswani et al., 2017 ] with hidden dimension of 128 and 4 heads in its attention. The positional embedding has length of 5, and GELU [ Hendrycks and Gimpel, 2016 ] and layer normalization [ Lei Ba et al., 2016 ] is used throughout the network. After the Transformer blocks, the output is fed into a layer normalization and a linear output layer to return logits. We use cross entropy loss to train the network and an Adam [ Kingma and Ba, 2014 ] with betas ( β 1 , β 2 ) = ( 0.9 , 0.98 ) (\beta_{1},\beta_{2})=(0.9,0.98)\, , a constant learning rate of 10 − 3 10^{-3}\, , batch size of 512 512\, , and linear learning rate warmup schedule over the first 10 iterations.
### B.2 MNIST
We train a three-layer MLP with hidden width of 200 and ReLU activations for the MNIST classification task [ Deng, 2012 ] . Under × 8 \times 8 larger weight initialization than Kaiming initialization [ He et al., 2015 ] , the network is known to exhibit the grokking phenomenon [ Liu et al., 2022b ] . The network receives flattened grayscale images of size 28 × 28 28\times 28 and outputs 10-dimensional logits to calculate MSE losses between one-hot encoded labels. We use the batch size of 200 200 and trained with an AdamW optimizer [ Loshchilov and Hutter, 2018 ] with a constant learning rate of 10 − 3 10^{-3} until 10 5 10^{5} training iterations. We use a smaller subset of 1000 images from training images to train the network in order to simulate overfitting environment. All the other hyperparameters are set by default.
### B.3 QM9
To demonstrate the effectiveness of our algorithm on a graph convolutional neural network, we use QM9 small molecules dataset [ Ruddigkeit et al., 2012 , Ramakrishnan et al., 2014 ] to estimate the isotropic polarizability. Our Graph ConvNet has two graph convolution layers with input channel of 11 (QM9 edge features), output channel of 16, and hidden channel of 32. Each graph convolution is followed by a ReLU. Each convolution layer consists of two linear layers with an internal ReLU activation with hidden channel of 32. The output of the Graph ConvNet is a global average pooling, followed by a two-layer MLP with a ReLU and hidden channel of 32. To simulate the overfitting environment, we use the first 100 samples from the data. The data is again randomly split into train and validation sets with 50:50 size ratio. We use batch size of 32, an AdamW optimizer [ Loshchilov and Hutter, 2018 ] with a constant learning rate of 10 − 3 10^{-3}\, . The network is initialized with weights × 3 \times 3 larger than that of Kaiming initialization [ He et al., 2015 ] and trained for 50k iterations.
### B.4 IMDb
For IMDb dataset [ Maas et al., 2011 ] , we use LSTM [ Hochreiter and Schmidhuber, 1997 ] with two layers, embedding dimension of 64, hidden dimension of 256, and vocabulary size of 1001, including the padding token. The network is followed by a single fully connected layer with output dimension of 1 with sigmoid activation to classify the positive/negative sentiment of each review string. The dataset was preprocessed by tokenizing the 1000 most frequent words from the review. The list of integer tokens are padded by zeros to form an array of reviews with the same length of 500. The network was trained by a binary cross entropy loss and an AdamW optimizer [ Loshchilov and Hutter, 2018 ] with learning rate of 3 × 10 − 4 3\times 10^{-4} and batch size of 50. We trained the model with the first 1000 rows from the dataset, split into train and validation sets with 75:25 size ratio. We stopped the training at 10k iterations as shown in Figure 11 .
## Appendix C Autograd Implementation
We have argued that our implementation of Algorithm 1 and 2 costs only a few additional lines of code. We demonstrate this by presenting the exact code we developed with the PyTorch [ Paszke et al., 2019 ] autograd package. The readers who are interested can also consult our official implementation 2 2 2 https://github.com/ironjr/grokfast .
Algorithms 1 and 2 are implemented as follows: ⬇