Grokking at the Edge of Numerical Stability

Authors: Lucas Prieto, Melih Barsbey, Pedro A. M. Mediano, Tolga Birdal
Primary area: learning theory
Keywords: grokking, deep learning, learning theory, floating point, scientific computation, generalization

Abstract:
Grokking, or sudden generalization that occurs after prolonged overfitting, is a surprising phenomenon that has challenged our understanding of deep learning. While a lot of progress has been made in understanding grokking, it is still not clear why generalization is delayed and why grokking often does not happen without regularization. In this work we argue that without regularization, grokking tasks push models to the edge of numerical stability, introducing floating point errors in the Softmax that we refer to as _Softmax Collapse_ (SC). We show that SC prevents grokking and that mitigating SC leads to grokking _without_ regularization. Investigating the root cause of SC, we find that beyond the point of overfitting, the gradients strongly align with what we call the _naïve loss minimization_ (NLM) direction. This component of the gradient does not change the predictions of the model but decreases the loss by scaling the logits, usually through the scaling of the weights along their current direction. We show that this scaling of the logits explains the delay in generalization characteristic of grokking, and eventually leads to SC, stopping learning altogether. To validate these hypotheses, we introduce two key contributions that mitigate the issues faced in grokking tasks: (i) $\mathrm{StableMax}$, a new activation function that prevents SC and enables grokking without regularization, and (ii) $\perp\mathrm{Grad}$, a training algorithm that leads to quick generalization in grokking tasks by preventing NLM altogether. These contributions provide new insights into grokking, shedding light on its delayed generalization, reliance on regularization, and the effectiveness of known grokking-inducing methods.

Decision: accept (poster)
Reviewer scores: 8, 5, 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1730669790727

thread_type:
official_review

<thread_object>
creator: reviewer_2l9f

<content>
summary:
This paper addresses the phenomenon of grokking in deep learning, where a model unexpectedly starts to generalize well after a period of severe overfitting. To combat this, the authors propose two primary methods: (1) a numerically stable StableMax activation function aimed at mitigating softmax collapse (SC), and (2) ⊥Grad (orthogonal gradient optimization) to counteract Naïve Loss Minimization (NLM), where gradients become aligned in a way that fails to drive effective learning. These methods are tested on a range of tasks, with a specific focus on scenarios where delayed generalization is a challenge.

strengths:
**Novel Approach to Softmax Collapse**: The introduction of the StableMax activation is an innovative approach to addressing potential issues with Softmax, providing a novel perspective on controlling model output stability.

**Orthogonal Gradients for Effective Training**: The concept of ⊥Grad is compelling and introduces a new approach to optimization that could be useful in a variety of training settings where standard gradient descent shows limitations.

**Comprehensive Analysis**: The paper provides detailed insights into the mechanics of grokking and includes empirical examples to illustrate the phenomenon.

weaknesses:
**Overemphasis on StableMax**: The proposed StableMax activation may be overengineered, as issues with softmax stability are often handled with simpler methods, like temperature scaling or loss smoothing. The authors could provide a comparison with such methods.
The practical necessity and advantage of StableMax might be limited in many standard applications.

**Underdeveloped Exploration of ⊥Grad**: Although ⊥Grad is promising, the current study lacks a thorough exploration of how it compares to more traditional optimization methods across a variety of tasks. Additional benchmarking could help clarify its broader applicability and effectiveness.

questions:
1. How does StableMax compare to standard techniques like loss smoothing? If StableMax’s advantage lies mainly in mitigating SC, could similar results be achieved with regularization techniques already in use? Some experiments with loss smoothing and temperature scaling to establish a fair comparison could be beneficial to the paper.

2. Is ⊥Grad beneficial for tasks where overfitting is not prominent? Understanding if this method generalizes to other contexts could add depth to the findings.

3. ⊥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? 
The paper would benefit from additional experiments in varied settings to benchmark ⊥Grad against standard optimizers and even on non grokking tasks. 

4. What is the computational overhead of implementing StableMax and ⊥Grad? Practical adoption would require insight into the trade-off between stability and additional computational complexity.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730680652625

thread_type:
official_review

<thread_object>
creator: reviewer_numh

<content>
summary:
This paper investigates why grokking happens and why it often does not happen without regularization. They show that lack of regularization in tasks where grokking is observed, can lead to overfitting, which is caused by floating point errors in the softmax. Mitigating this softmax collapse (SC) leads to grokking. Further, they show the root cause of SC to be the alignment of gradients with a naive loss minimization (NLM) direction, which does not improve generalization, even though it reduces train loss via scaling. Finally, they propose a StableMax activation that prevents SC and a training algorithm $\perp$ Grad that leads to faster generalization, by amplifying gradients orthogonal to the NLM direction.

strengths:
The paper advances our understanding of why generalization or grokking may not happen without regularization and proposes ways to induce grokking or faster generalization via StableMax and $\perp$ Grad.

The paper is well-written and easy to understand. The claims are well supported by the presented results.

weaknesses:
There are no major weaknesses. Some minor concerns are as follows.

- The phase ‘grokking tasks’ can be reworded as ‘tasks where grokking is commonly observed’, since there is no formal definition of grokking tasks.
- The readability can be improved by better placement of the figures (same page as the corresponding discussion).
- In lines 302-304, the digits 1- can be replaces with i) for better readability.

questions:
- How do other factors such as initialization scale, learning rate, etc. interact with grokking?
- Do the results shown in Fig. 4 correspond to one the settings shown in Fig. 2?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730696177184

thread_type:
official_review

<thread_object>
creator: reviewer_yrj9

<content>
summary:
The paper studies grokking. It found that when the loss function is cross-entropy loss, the floating point errors caused by softmax prevents the model from generalization if trained without any regularization. They proposed a modified loss function and a modified optimizer that  enable the model to generalize without regularization.

strengths:
This work studies grokking from an interesting view and provided an explanation for why the model is trapped in an overfitting regime.

weaknesses:
- The application of the proposed method is limited. It only turns from overfitting without generalization to generalization when there is no regularization and the original loss is cross-entropy loss. Is ⊥AdamW still better than AdamW with non-zero weight decay? E.g. can you compare ⊥AdamW against AdamW with various non-zero weight decay values for transformers and MLP on modular addition task?
- The proposed method is only tested on some toy models such as two-layer neural networks and one-layer transformer. It would be more convincing to apply the method on deeper transformers, e.g. the same architecture as GPT2-small.
- Sec 4.1 compares one-hot encoding with binary and scalar representation to show overfitting is harder, which is not fair: 
    - 1. the one-hot vector does not have any information of the integers, but binary and scalar vectors do contain some info. 
    - 2. the model trained on binary and scalar vectors has fewer parameters, making overfitting harder. 
A fair comparison would be using randomly generated d-dim vectors on a uniform sphere as the low-dim representation, and using a model with larger width to make sure the number of params is roughly the same.

questions:
- What’s the difference between absorption error and Roundoff error?

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

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

- [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?

[1] Progress measures for grokking via mechanistic interpretability

[2] The Slingshot Mechanism: An Empirical Study of Adaptive Optimizers and the Grokking Phenomenon
</content>
</thread_object>

</thread>
<thread>
creation date:
1730897082543

thread_type:
official_review

<thread_object>
creator: reviewer_sm8g

<content>
summary:
The paper investigates the phenomenon of grokking, where deep learning models suddenly generalize after a period of overfitting, and it proposes new mechanisms that explain why generalization is often delayed. The authors identify Softmax Collapse, a numerical instability in Softmax, as a key factor preventing grokking in models trained without Weight Decay. To address this, they introduce "StableMax," a modified Softmax function that reduces numerical errors, and "⊥Grad," an optimizer that excludes the naive loss minimization direction.  The latter refers to an optimization path that reduces the loss by scaling logits without changing model predictions and potentially leading to numerical errors.

strengths:
- The paper is well-written and easy to follow.
- The research problem is conceptually well-motivated and aligns with current challenges in understanding grokking.
- The hypothesis presented, that Softmax Collapse inhibits grokking, is interesting.

weaknesses:
1. **Unclear Justification for Proposed  Adjustments Over Established Regularization Techniques** While the paper suggests several adjustments to address generalization delay, it remains unclear if and why these modifications should be preferred over established techniques such as weight decay, particularly at larger scales. A comparative analysis with standard approaches would strengthen the argument for these adjustments.

2. **StableMax Analysis:** The introduction of StableMax is a substantial change to the training paradigm. However, the paper lacks an in-depth analysis of its properties, such as the impact on the optimization landscape and the types of solutions this modified model converges towards. Given the importance of StableMax, its introduction should be supported by either extensive experimental validation or a rigorous theoretical justification.

3. **Optimizer in Section 5.1:** Similar to StableMax, the optimizer proposed in Section 5.1 would benefit from additional empirical or theoretical analysis. Specifically, how it compares with other optimizers in terms of convergence behavior and stability beyond the context of grokking tasks.

4. **Insufficient Evidence for Naive Loss Minimization:** The evidence provided for the claim of naive loss minimization is limited. Existing literature, such as that referenced work of Lyu and Li (2019)​[1] (in Appendix K.1 ), shows that even after achieving 100% training accuracy, gradient descent with cross-entropy loss implicitly performs margin maximization. This result seems to contradict the authors’ hypothesis that the optimization follows a direction of naive loss minimization, where the function remains unchanged across the entire input space (as described in Equation 9).  Further evidence is needed to validate naive loss minimization from this established understanding of implicit bias in gradient descent.

questions:
1. Could the authors further discuss naive loss minimization in relation to the max-margin implicit bias as described in [1]? In particular the claim that the direction of the optimization trajectory is the one that does not change the prediction for all data points in the input space (Equation 9). 


2. Could the authors clarify why the proposed adjustments should be prioritized over traditional methods like weight decay?


3. Can the authors elaborate on how the proposed optimizer relates to [2]?


4. Is there a connection between the projected optimizer in this paper and the rotational equilibrium update discussed in [3] as an effect of weight decay?


5. For Figure 2, could the authors specify the task on which the models are trained?


6. In Figure 6, do the authors have experiments demonstrating the impact of different weight decay values on the delay in generalization?

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


[2] Heo, Byeongho, et al. "Adamp: Slowing down the slowdown for momentum optimizers on scale-invariant weights." _arXiv preprint arXiv:2006.08217_ (2020).


[3] Kosson, Atli, Bettina Messmer, and Martin Jaggi. "Rotational equilibrium: How weight decay balances learning across neural networks." _arXiv preprint arXiv:2305.17212_ (2023).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730965286034

thread_type:
official_review

<thread_object>
creator: reviewer_mcra

<content>
summary:
This paper uncovers the critical role of numerical stability in grokking, demonstrating that Softmax Collapse  (SC) prevents grokking and that mitigating SC leads to grokking without regularization.

strengths:
- The authors make a novel discovery regarding the crucial role of numerical stability in grokking and identify the Softmax Collapse (SC) phenomenon.
- They reveal that SC results from naive loss minimization (NLM) of cross-entropy loss.
- Based on these theoretical insights, the authors propose two methods—StableMax and $\perp$Grad—to mitigate issues faced in grokking tasks.

weaknesses:
- **Softmax Collapse (SC).** 
  - The current definitions in Section 3.1 is somewhat abstract and could be more practically grounded. For example, when using float16, what is the $\epsilon$ in Definition 1, and what specific difference between the two logits in Definition 3 would lead to SC?
  - In transformers, softmax is applied in all self-attention layers, not just the last layer. Does SC also occur in these layers?

- **StableMax.**
  - I am concerned that StableMax may not be suitable for the softmax in attention layers of transformers. Many experiments show that attention matrices in transformers are highly sparse, but the *slower growth rate* of StableMax (compared to softmax) could hinder its ability to model such sparsity.

- **Naive Loss Minimization (NLM).**
  - It is worth noting that this emperical observation has been proven in related works, such as in linear models [1], deep linear networks [2], and homogeneous networks [3]. As training progresses, the "angular velocity" becomes much slower than the growth of the norm, and the directional convergence holds: $\left<-\frac{\nabla L(\theta_t)}{\|\|\nabla L(\theta_t)\|\|},\frac{\theta_t}{\|\|\theta_t\|\|}\right>\to1$.
  - It seems that this explanation holds only for cross-entropy loss. For squared loss, NLM does not appear to hold.

- **$\perp$Grad.** 
  - Based on my understanding, the motivation behind $\perp$Grad is to amplify the angular velocity of GD, similar to the effects of normalization layers and weight decay. Additionally, related work [4] also explores how to amplify the angular velocity of GD.
  - Adaptive algorithms like Adam may not encounter the same NLM problem as GD late in training. Would it be effective to use $\perp$Grad+Adam to accelarate the grokking process of Adam?


References:

[1] Ji and Telgarsky. Risk and parameter convergence of logistic regression. COLT 2019.

[2] Ji and Telgarsky. Gradient descent aligns the layers of deep linear networks. ICLR 2019.

[3] Ji and Telgarsky. Directional convergence and alignment in deep learning. NeurIPS 2020.

[4] Wang et al. Achieving Margin Maximization Exponentially Fast via Progressive Norm Rescaling. ICML 2024.

questions:
Refer to Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732294179952

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank all reviewers for their constructive feedback. We are happy to see that the reviewers appreciate the **critical role** [Mcra] that the dynamics uncovered in this paper play in grokking and believe that it **advances our understanding** [NUmH] of this phenomenon, with other reviewers agreeing that the work is **interesting** and **well motivated** [sm8G, yRj9] and the proposed solutions **innovative** and **compelling** [2L9F]. We now address the general concerns.

**Evaluation in more realistic settings:** While the StableMax and $\perp$Grad are mainly proposed as interventions to validate our novel hypotheses and advance our understanding of grokking and generalization, we appreciate the reviewers' interest in seeing how these methods would perform in more realistic settings. Accordingly, we have added experiments using ResNet18 on Imagenet-1k, CIFAR100 and CIFAR10, as well as GPT2-Small on Wikitext-103,the results in the table below, show that outside of the extreme overfitting setting they were designed for, $\perp$Grad and StableMax cross-entropy achieve performances which are comparable with our baseline for language modeling and image classification, with $\perp$Grad outperforming our baseline for CIFAR100. StableMax is designed to be as similar as possible to softmax while preventing softmax collapse, therefore, it seems to behave very similarly to softmax in most cases, however, we notice that when StableMax is applied over a large number of classes or tokens (50,000 for GPT2) it does seem to degrade performance from 59.58\% to 52.14\% top-5 test accuracy and a less significant decrease is also observed for Imagenet-1k (69.76\% to 65.76\%). We hypothesize that the difference between softmax and StableMax becomes more significant when applied over large dimensions, where the probabilities assigned to the large number of incorrect classes does not decrease as quickly with StableMax, interfering with the signal from the correct label. In these settings increasing the temperature of StableMax or other interventions might be needed to use StableMax over very large dimensions. Overall, we find that \textit{i}) in cases where numerical instability is a significant risk, StableMax and $\perp$Grad bring substantial benefits from avoiding such instability; and \textit{ii}) when there is no numerical instability, in most cases StableMax and $\perp$Grad have no negative effect (except StableMax applied over large dimensions) -- as expected from our theoretical motivation. The table below is included in Appendix G of the revised paper, and reports test accuracies which are averaged across 5 seeds and show the standard deviation for CIFAR10 and CIFAR100. We are currently running different seeds for Imagenet-1k and WikiText-103 to be included in the final version.


| **Method**             |    **CIFAR10**        |    **CIFAR100**       |       **ImageNet-1k**.   |    **WikiText-103 (Top-5)** |
|------------------------|--------------------|--------------------|-----------------|--------------------------|
| Softmax CE             | 87.17\% ± 0.2      | 59.98\% ± 0.4       |     69.76\%          |    59.58\%                  |
| StableMax CE           | 87.01\% ± 0.2      | 60.63\% ± 0.4       |     65.76\%          |   52.14\%                  |
| $\perp$Grad              | 87.22\% ± 0.2      | 62.69\% ± 0.1       |    68.9\%           |    59.53\%                  |
| StableMax Attention| --                | --                 | --              |    58.37\%                  |

**$\perp$Grad and weight decay:** Following the helpful suggestion from the reviewers, we have added a more detaield comparison of weight decay and $\perp$Grad. Our results show that higher weight decay can anticipate grokking, but only up to a point, beyond which it prevents generalization. We then compare a tuned value of weight decay, chosen to induce grokking as early as possible and see that it is still marginally slower than $\perp$Grad which does not require any hyper-parameters. Additionally, we observe that AdamW with weigh decay overfits very quickly, leaving a larger delay until generalization, while with $\perp$Grad, train and test accuracies reach 100\% almost at the same time. These results have been added to the paper in Figure 15 (Appendix E).


**Related literature:**
We also thank the reviewers for pointing out some relevant literature, particularly, existing works related to $\perp$Grad and NLM. We have updated the paper to include a short discussion of these methods and their differences from $\perp$Grad. As suggested by reviewer 2L9F, we have also included a comparison of StableMax to label smoothing and temperature scaling in the Appendix.  

More detailed descriptions of these changes are provided in the answers to individual reviewers.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733168214479

thread_type:
public_comment

<thread_object>
creator: ~vimal_thilak1

<content>
title:
Models used in paper

comment:
I read the paper with great interest and think that this paper is good. I am an author that contributed to [Slingshot Effect]((https://openreview.net/forum?id=OZbn8ULouY) work that is referred to by the authors in this paper. I am curious to know whether the authors used a 2 layer Transformer with the datasets and setup described by Power et al. in their original grokking paper? Looking at lines 519-524, I am curious to understand whether the authors considered some of the observations made by us while training Transformers without with Adam (and no weight decay)? One attempt to resolve instabilities during optimization was to increase epsilon but this resulted in no generalization (and consequently no grokking) which is undesirable.  

We also have a discussion on scaling weights and activations of the last layer in an appendix (C.2) in [our paper](https://openreview.net/forum?id=OZbn8ULouY). If the authors find our observations relevant then would they perhaps consider discussing these specific observations as related work in a future revision?

Best,

Vimal
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper aims to explain why grokking appears and disappears (without regularization) for several existing methods to induce grokking. It identifies the crucial role of numerical stability in grokking and identify the Softmax Collapse (SC) phenomenon. Accordingly, the authors develop a simple modification of the Softmax to avoid floating point errors, which will induce grokking. The extensive experiments can support this claim.

It's true that the training dynamics is missing, particularly its relationship to grokking and generalization, which indeed decreases the significance of this work. But anyway, this paper is good in general, and I recommend to accept.