Grokking at the Edge of Numerical Stability
grokking;deep learning;learning theory;floating point;scientific computation;generalization
Yes
No ethics review needed.
Refer to Weaknesses.
- 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.
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.
- **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.
Yes
No ethics review needed.
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).
- 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.
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.
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.
Yes
No ethics review needed.
- 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
This work studies grokking from an interesting view and provided an explanation for why the model is trapped in an overfitting regime.
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.
- 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.
Yes
No ethics review needed.
- 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?
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.
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.
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.
Yes
No ethics review needed.
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.
**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.
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.
**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.
We study the grokking phenomenon and identify issues with numerical stability in the cross-entropy loss. We develop new approaches to induce grokking without regularization.
@inproceedings{
anonymous2024grokking,
title={Grokking at the Edge of Numerical Stability},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=TvfkSyHZRA},
note={under review}
}
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.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
grokking deep learning learning theory floating point scientific computation generalization
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/3ac0ffded48648f3688ceecb641f7c077f87a60c.pdf
learning theory
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/f16cf9b846d080def64a3556dd84f9b5727cf26e.zip
Grokking at the Edge of Numerical Stability
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission