Grokking at the Edge of Numerical Stability

We thank the reviewer for this positive feedback. We have made some changes to address the minor concerns.

Q1: How do other factors such as initialization scale, learning rate, etc. interact with grokking?

Our observations align with the literature in observing that learning rate can speed up grokking up to a point and that scaling the weights can delay grokking. We additionally observe in Appendix D.2 that scaling the weights also scales the logits, leading to faster Softmax Collapse and explaining how delayed generalization can be induced in MNIST by scaling the weights and using cross-entropy loss.

Q2: Do the results shown in Fig. 4 correspond to one the settings shown in Fig. 2

Yes, thank you for pointing this out, we have updated the caption fro Fig 4 to make this clearer.

We thank the reviewer for the constructive feedback. We provide answers to the reviewers questions below:

Comparing $\perp$AdamW to AdamW with weight decay: Thank you for the suggestion. In addition to the results with weight decay we have provided in Figure 6, we have added a comparison of $\perp$AdamW against AdamW with the weight decay parameter selected to induce grokking as early as possible. The results show that weight decay can anticipate grokking up to a point, beyond which weight decay prevents generalization. Even without any tunable hyper-parameters, $\perp$Grad leads to faster grokking than any value of weight decay. These results have been included in Appendix E.

Evaluation in more realistic settings: We agree with the reviewer that it is interesting to observe the behavior of our interventions in the beyond-the-toy settings. We take this suggestion into account and added experiments with more realistic models like Resnet18 and GPT2-Small on datasets like Imagenet and GPT2-Small. These results are included and discussed under Evaluation in more realistic settings in the general response and have been added to the paper in Appendix G.

One hot vs binary encoding: We thank the reviewer for the insightful comment. We agree that the previous setup was not a fair comparison as the binary and scalar representations introduce information about the structure of the integers making the task easier as well as making overfitting harder. As a remedy, we have updated this setup such that the inputs involve random permutations applied to the integers before converting them into binary representations and concatenating them, thereby removing any information about the integers being encoded. The results indicate that reducing the dimension of the input data leads to direct generalization, even if this new representation does not introduce any additional information about the integers and we adjust the number of parameters (Increasing the hidden size from 200 to 280 for a roughly equal parameter count when decreasing the input size from 226 to 14). These results have been updated in Figure 4.

Absorption vs roundoff errors: A Roundoff error is any difference between the result of a sum with floating point numbers and the true result of that sum. Absorption is a special case of Roundoff error that happens when adding a very large and a very small floating point number such that the smaller number is "absorbed".

Q2: Why setting the gradient of correctly predicted (overfitted) samples to zero prevents generalization? Is there an intuitive explanation?

Zero-gradients of overfitted samples: While the training accuracy is 100%, cross-entropy loss should never be 0 in the absence of floating point errors. It has been proven in [3] that homogeneous models perform normalized margin maximization beyond the point of 100% training accuracy and we believe this is the process that leads to grokking unless it is stopped by Softmax Collapse (SC).

Q3: In Figure 1c, can grokking still occur after careful hyper parameters(initialization scale, learning rate) tuning?

Grokking and hyper-parameter tuning: We do not believe that grokking can be induced in this setting without introducing an intervention that prevents SC. In Figure 1c we show that grokking does not occur even after training for 80,000 epochs in a setting where generalization happens in less than 200 epochs with $\perp$Grad. Both of these use AdamW with regular initialization scale and the same very common hyper-parameters (lr=0.001 and default AdamW with weight decay set to 0). Additionally, Figure 2a shows that in this same setting, the gradients coming from the correct class for most of the samples become 0, preventing generalization beyond that point.

Q4: [1] found that one-layer transformer trained with float32 precision has loss spikes and use float64 for more stable training. This implies that the slingshot mechanism [2] may be caused by numeric stability. Can softmax collapse explain slingshot?

Can our findings explain slingshot mechanism? We thank the reviewer for bringing this interesting insight. Indeed, based on the observations from [1], we believe softmax collapse plays an important role in the creation of slingshots, when this is combined with the second order moments in Adam (which were found to be necessary for slignshots in [2]). We hypothesize that zero gradients caused by SC lead to very small second order moments resulting in slingshots when these weights get non-zero gradients, but we do not have enough experiments on this front to be sure. This is briefly mentioned under `numerical instability in deep learning' in Related Work.

We thank the reviewer for the constructive feedback. Below, we address the reviewers questions and concerns. While we appreciate the reviewer's detailed feedback, we want to highlight that we consider the key contributions of this paper to be more about convincingly diagnosing the problems faced in grokking than about the solutions proposed. It is commonly believed in the literature that grokking with cross-entropy loss requires weight decay (or some regularization), and several explanations have been proposed as to why this is the case. In contrast, we show that gradient descent leads to grokking by default (without weight decay), unless this is stopped by floating point errors and explain why these floating point errors appear in grokking (due to NLM). To show this, we introduce StableMax as the simplest change to the Softmax that keeps it from collapse and $\perp$Grad as the most direct intervention to prevent naïve loss minimization (NLM).

Comparing StableMax with label smoothing and temperature scaling: Thank you for the suggestion. We performed additional experiments to compare label smoothing, temperature scaling and StableMax. We found that even with very large temperature (T=1e+5), models eventually face softmax collapse (SC), and do so very quickly when using optimizers with momentum like Adam. Label smoothing does prevent SC and leads to some generalization, although label smoothing introduces other inductive biases that lead to imperfect generalization and qualitatively different behavior than the one commonly observed in grokking (Figure 16). We believe StableMax provides a targeted intervention, where a single change to the Softmax prevents SC and is the difference between grokking and overfitting. These results have been added to the paper in Figure 16 (Appendix F).

Evaluating $\perp$Grad in realistic settings: We have conducted additional experiments comparing $\perp$Grad to other methods and show that it is competitive with our baselines in Imagenet and WikiText, and outperforms them in CIFAR100. However, as suggested by the reviewer the main benefits of this method seem to come in settings wehre overfitting is prominent. These results are included in Appendix G and we also show the full results for $\perp$Grad in the table below.

Q3: $\perp$Grad prevents scaling in the logits; is there any impact on margin and generalization, or is it compensated because the norm of weights remains low?

We thank the reviewer for this interesting question. When studying margin maximization in neural networks, the normalized margin is commonly studied for the precise reason that margin maximization through scaling all the logits does not lead to generalization [1]. Therefore, the fact that $\perp$Grad prevents scaling in the logits should not (and empirically does not) affect generalization.

Q4: What is the computational overhead of implementing StableMax and $\perp$Grad? Practical adoption would require insight into the trade-off between stability and additional computational complexity.

We thank the reviewer for pointing this out, while StableMax is as fast as Softmax and introduces no additional computational overhead, the gradient projection in $\perp$Grad does account for around 38% of the training time, in the case of a GPT2-Small trained on WikiText-103 and 34% for a Resnet18 trained on Imagenet. While we mainly use $\perp$Grad to show to highlight the role of NLM in grokking, if one wants to apply this to more realistic settings, this should be a consideration.

We hope answers and additional experiments have addressed the reviewers concerns, and would allow them to consider updating their score.

[1] Lyu, Kaifeng, and Jian Li. "Gradient descent maximizes the margin of homogeneous neural networks." arXiv preprint arXiv:1906.05890 (2019).

We thank the reviewer for the constructive feedback and provide detailed answers to their questions below.

Connection to Lyu and Li (2019): The connection to the work from Lyu and Li ([1]) is very relevant. While we define the NLM direction as not changing the predictions for any data points (Equation 9), we argue that in practice, models strongly "align" with the NLM direction with large cosine similarities above 0.9 in the MLP case (Figure 5), but they do not exactly follow the NLM direction. The part of the gradient that is orthogonal to the NLM direction does increase the margin as predicted by Lyu and Li and validated by the success of $\perp$Grad which only preserves this part of the gradient. Lyu and Li also run into numerical issues in prolonged training beyond 100% accuracy (Appendix L.2). In their case these issues can be fixed by using float64, but we show that in cases of grokking this is not enough, and diagnose the particular kind of floating point issues to be Softmax Collapse. To take the connection to this paper further, we believe the max-margin implicit bias described in Lyu and Li (2019) is the reason we should observe grokking by default. The reason this is not observed in the literature when using cross-entropy loss is that Softmax Collapse stops this margin maximization process. While Lyu and Li prove this max-margin property for gradient flow and gradient descent, the proofs do not account for limited floating point precision. We have updated the paper to make it clearer that our findings are in line with this established implicit bias of GD (Section 4.2).

Comparison to weight decay StableMax and $\perp$Grad are mainly presented as targeted interventions to validate our hypothesis that floating point errors caused by NLM inhibit grokking and are not intended to replace weight decay. However, we have added results showing that weight decay only anticipates grokking up to a point (weight decay=6 in our case) beyond which it prevents generalization (Figure 15a). We compare an MLP trained on modular addition using the tuned weight decay value and show that $\perp$Grad still leads to faster grokking than weight decay without requiring tuning of the weight decay parameter (Appendix E, Figure 15b).

Additional literature: We thank the reviewer for pointing out this relevant literature. The rotational equilibrium described in [3] could explain why weight decay behaves like $\perp$Grad, but only after a delay. In the terms of [3], the delayed generalization in grokking could correspond to the transient phase, with generalization starting when the model reaches rotational equilibrium. More work would be needed to make this connection rigorous. While the optimizer in [2] is designed for scale invariant weights and does not prevent scaling for regular weights, some of the insights about how to handle momentum in these settings could be incorporated to improve the performance of $\perp$Grad when combined with Adam ($\perp$Adam). We have updated Section 5.1 to mention this and cite these works.

Question: For Figure 2, could the authors specify the task on which the models are trained?

Answer: We thank the review for pointing this out, we have updated the figure caption to mention that these results are from modular addition.

General note: While this paper's main goal is to gain insight into the grokking, we agree that an exploration of the broader applicability of these methods can strengthen the paper. As such, we have conducted experiments (results in the table under the General Response) using GPT2-Small, and Resnet18 that show that our methods behave similarly to our baseline in the context of language modeling (WikiText-103) and image classification (Imagenet-1k, CIFAR100, CIFAR10), with the exception of StableMax which behaves almost identically to softmax when predicting across a limited number of classes (CIFAR100) but decreases performance when predicting over 50,000 tokens (WikiText-103). Fot a more detailed discussion of these new results we point the reviewer to the general response. These results are also included and discussed in the updated paper under Appendix G.

We thank the reviewer for the constructive feedback. Below, we address the reviewers questions and suggestions.

Clearer definition of absorption errors: We thank the reviewer for pointing this out. Type float32 uses 24 significand bits, meaning that adding any number smaller than $2^{-24}$ to 1 ($1=2^0$) will leave 1 unchanged. $2^{-24}$ is called the machine epsilon for float32. We have now clarified this in the main paper (Section 3.1).

StableMax in the attention mechanims of a transformer: Although StableMax is designed as a minimalist intervention to validate our hypothesis that preventing SC leads to grokking, we share the interest of the reviewer in seeing if this could be included in the attention mechanism of a transformer. We have added experiments using the GPT2-Small architecture which show that using StableMax instead of Softmax in the attention mechanism has a small negative impact on performance, from 59.58% to 58.37% top-5 accuracy. This suggests that the sparsity induced by softmax may be beneficial in the attention mechanims, as suggested by the reviewer, although the effect is small and more extensive experimentation is needed to conclude this definitively. These results have been added to Appendix G.

Additional related work: We agree that NLM is closely related to the phenomena discussed in these works ([1], [2], [3]) as well as [5], and have included a short discussion of this as well as the relevant citations in Section 4.2 . While the convergence of the gradients to the direction of the weights is not a new observation, we highlight that this leads to the scaling of the logits, which can become a problem when we take limited floating point precision into account.

Relevance for MSE loss: The reviewer is correct in highlighting that the phenomena we observe do not apply to MSE loss. As discussed in Section 5.2, this actually explains why grokking without regularization has been observed with MSE loss but not with cross-entropy, which was a previously unexplained phenomenon.

Concerns about $\perp$Grad: We agree with the reviewer that $\perp$Grad does amplify the angular velocity and is related to methods like [4], which uses progressive rescaling to address a limitation of normalized gradient descent and actually use a projection similar to $\perp$Grad as a way to stabilize their method. We have updated the paper to mention this in Section 5.1 . Regarding the reviewers question about adaptive algorithms, Adam does seem to face the same NLM and SC problem as SGD. As shown in Figure 1, AdamW without weight decay does not generalize and faces SC (Figure 2a), in settings where preventing SC does lead to grokking. Figure 1a also shows that $\perp$Grad+Adam (termed $\perp$AdamW in the paper) does accelerate grokking, as suggested by the reviewer.

[5] Lyu, Kaifeng, and Jian Li. "Gradient descent maximizes the margin of homogeneous neural networks." arXiv preprint arXiv:1906.05890 (2019).

We would like to inform the reviewers that, as discussed in our previous response, we have now updated the results table (Table 1 in the paper) to report the mean accuracy and standard deviation across different seeds for Imagenet-1k and WikiText-103. Below is the updated table:

For the methods introduced in this paper, we report test accuracies with standard deviations across five seeds for the CIFAR datasets and three seeds for ImageNet-1k and WikiText-103. We report Top-5 accuracy in the case of WikiText-103.