Grokfast: Accelerated Grokking by Amplifying Slow Gradients — part 4/4

Cleaned-text code points [47384, 51875). Bibliography omitted here; source.md includes it.

This helper method can be applied to any optimization framework involving the autograd package by inserting a single line between the calculation of the gradients and the optimizer call as follows:

Note that line 2 and line 14 in the code above are the only modification we made.

## Appendix D Time and Memory Requirements

This section delivers additional demonstration of the efficiency of our Grokfast algorithm. As we have argued in Section 2.4 and Section 3 , the additional computational burden from our augmentation is compensated by the larger-scale acceleration of the delayed generalization. The additional cost of VRAM memory is also negligible compared the baseline. The time and the memory requirements in the Tables 6 through 6 are measured with a single GTX 1080 Ti GPU.

## Appendix E More Visualization

We finally provide more visualization in addition to Section 5 in order to understand the training dynamics under our Grokfast algorithm. Figure 13 shows five more runs from the same experiments in Figure 2 and 12 with different seeds. We saved all the 423k parameters of the Transformer decoder [ Vaswani et al., 2017 ] from every training iteration, likewise in Figure 12 . The parameters of a model checkpoint from each run at each iteration are reshaped into a single long vector. The vectorized parameters are then normalized by subtracting them by the model’s initialized weights. This way, we can align the trajectories by centering the initial states (state A ) of all the experiments at the origin. The sequence of parameter differences of the ten runs from the two algorithms, i.e., baseline and Grokfast , forms a tensor of shape ((Number of Runs) ⋅ \cdot (Number of Iterations)) × \times (Number of Parameters) = ( (Number of Sampled Iterations) × 422784 ) (\textrm{(Number of Sampled Iterations)}\times 422784)\, . From this we perform the PCA to obtain the projection matrix of shape 422784 × 2 422784\times 2\, . This matrix projects the parameter differences from each of the model checkpoint onto the two most salient directions of variations. We mark the initialization state (state A ), the overfitting state (state B , 500 iterations), and the generalization state (state C ) from each run in the two-dimensional plot. The results are Figure 13 .

We first notice that the overfitting states (state B ) from each of the two optimizers are clearly different. The baseline algorithm without Grokfast reaches the overfitting states (state B ), which are relatively nearer to the initialization states (state A ) than those of Grokfast algorithm. However, as soon as the model overfits, the weights continue to deviate far from the points where overfitting first occured (state B ). As a result, the final generalization (state C ) happens much far away from the initialized weights (state A ). It is notable that the five generalization states (state C ) from different instances of the baseline optimizer vary significantly. The difference between the states C of the baseline is much larger than that of the baseline’s overfitting states (state B ) and that of the states B and C from our Grokfast algorithm. In contrast, the training dynamics from Grokfast results in a distinct set of trajectories that lead to the generalization states (state C ) much closer to the initial weights. Moreover, difference within the trajectories from Grokfast is much smaller than that of the baseline algorithms. This conclusion is also verifiable from Figure and Table 2 in a quantitative manner.

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .