Grokking at the Edge of Numerical Stability
Authors: Lucas Prieto
Melih Barsbey
Pedro A. M. Mediano
Tolga Birdal
Venue: ICLR 2025 Poster (2025)
Keywords: grokking
deep learning
learning theory
floating point
scientific computation
generalization
Primary area: learning theory
TL;DR: We study the grokking phenomenon and identify issues with numerical stability in the cross-entropy loss. We introduce methods to prevent these issues, leading to grokking without weight decay.

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.

Published as a conference paper at ICLR 2025
GROKKING AT THE EDGE OF NUMERICAL STABILITY
Lucas Prieto, Melih Barsbey, Pedro A.M. Mediano", Tolga Birdal* Department of Computing Imperial College London
ABSTRACT
Grokking, sudden generalization that occurs after prolonged overfitting, is a surprising phenomenon that has challenged our understanding deep learn- ing. While lot of progress has been made in understanding grokking, it is delayed and why grokking often does not
still not clear why generalization happen without regularization. In this work argue that without regulariza- tion, grokking tasks push models to the edge of numerical stability, introducing floating point errors in the Softmax that we refer We show that SC prevents grokking and that mitigating SC leads to grokking without regularization. Investigating the root cause SC, we find that be- yond the point of overfitting, the gradients strongly align with what we call the naive 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 direc- tion. We show that this scaling of the logits explains the delay in generaliza- tion 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) StableMax, a new activa- tion function that prevents SC and enables grokking without regularization, and (ii) 1 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 regu- larization, and the effectiveness of known grokking grokking-inducing methods. Code for this paper can be found at: https:://itthhhi
Softmax Collapse (SC).
com/
grokking -aanng-ttttt at- tet-the-edde of numerical stability
INTRODUCTION
Deep learning has been transformative for a variety fields such as natural language process- ing (Devlin et al. 2019), computer vision (Krizhevsky al. 2012), geometry processing (Qi et al. 2017). and 3D vision (Deng et al.) 2018). This rapid proliferation has brought with surprising phenomena that defy the predictions of classical statistical learning theory.
In this paper we explore one such recently observed phenomenon known as grokking. first described by Power al. (2022) as sudden and unexpected generalization occurring after prolonged overfit- ting. Although predominantly studied in algorithmic tasks like modular addition or multiplication, recent findings suggest that grokking may a more pervasive phenomenon, also manifesting in
more complex tasks involving vision and language (Lv et al. 2024| Humayun et al. 2024).
Prior research has consistently observed grokking in settings that involve some form of regulariza- tion, such as weight decay (Barak et al.) 2022 Power et al. 2022 Nanda et al. 2023). This pattern has motivated investigations into the implicit blases introduced by weight decay, suggesting it may be critical to triggering delayed generalization. For instance, [Lu et al. (2023a) argued that weight norms need to be a narrow range or "Goldilocks Zone" for generalization. Similarly, Varma al. (203) highlighted weight efficiency of generalizing solutions, and Nanda et al.) (2023) argued that weight decay favors simpler, more generalizable solutions. However, recent works have argued that regularization may not be necessary for grokking, at least on shallow networks with Mean Squared
"Joint senior authors, equal contribution

---

Published as a conference paper at ICLR 2025
100
100
VC
100
due to logit-scaling NLM (Sec. 3)
-0
due numerical ical errors SC (Sec 2)
8
Triiing BOURO
[Scc 4.2)
(Sec. 2,4)
250 Epo
300
200 392
JAdamW StabloMax
13c
50. *0x Mex spech
Adoon + StahleMax
2pech
Generalization
Grokking
Overfitting
Figure I: Our contributions demonstrated through results obtained in addition modulo 113 task. We show that the delay in generalization induced by NLM can be reversed using the proposed LAdamW ((a) and (b)) and that the numerical errors that lead to overfitting instead of grokking be avoided by using the proposed StableMax ((b) and (c)).
Error (MSE) loss (Kumar et al.) 2024 Lyu et al. 2024 Gromov 2023). These works tie grokking to transition from lazy training (Chizat et al., 2018) feature learning. Despite this ongoing work, several aspects in this framing of grokking remain unclear. These include why grokking tasks induce lazy training and why weight decay is often needed to enter the feature learning regime when using deeper models or cross-entropy (CE) loss.
Here we propose a novel account of grokking, outlined in Fig. that explains several of the main unanswered questions in the grokking literature. We start by showing that without regularization, grokking is prevented by absorption errors in the Softmax, which we call Softmax Collapse (SC). These errors result in zero terms in the gradient and put an end to learning, sometimes before any progress made the test performance, resulting in complete overfitting (Fig. I c). We then argue that SC is caused by what call Naive Loss Minimization (NLM), the gradient becomes aligned with a direction that corresponds to scaling up the logits by constant. While scaling up all the logits does not change the model predictions, it does reduce the CE loss for network that has reached 100% training accuracy, with the downside that this eventually leads to numerical errors in Softmax. Our findings provide explanations for several key aspects of grokking, including (i) the delayed onset of generalization, (ii) why grokking is often absent without regularization, and (iii) why existing methods designed to induce grokking are effective.
To validate our hypothesis that SC is responsible for the absence of grokking without regularization, we introduce StableMax as a more numerically stable replacement Softmax in CE loss. This simple change takes models from complete overfitting to grokking (Fig. [ с to b) without regular- ization, in settings where is normally not observed without it. Similarly, we validate that NLM is responsible for delaying generalization (Fig. [ a to b) and leading to SC by introducing new optimizer 1 Grad, which only preserves the part of the gradient that orthogonal to the NLM di- rection. By doing this, 1 Grad quickly leads to generalization without the initial overfitting phase that defines grokking (Fig. [ a).
Our primary contributions are as follows:
We observe that cases of overfitting without grokking are due to floating point errors caused by extreme values in the Softmax function, which we term Softmax Collapse (SC; Sec. 3
We show that interventions to avoid SC, like greater floating point precision or a new, numerically stable version of Softmax (StableMax). cause grokking in settings where it was previously absent without regularization (Sec.B.3)
We observe that models move towards SC because overfitting and cross-entropy loss push the model in a direction of uncontrolled logit growth, which we refer to as Naïve Loss Minimization
(NLM: Sec.
We demonstrate that NLM can be avoided through a novel optimizer, _Grad, which removes the delay in generalization (Sec.
SETUP
2.1 DATASETS
We show our findings on the most commonly studied grokking datasets, outlined in this section.
2

---

Published as a conference paper at ICLR 2025
I. Modular arithmetic. The main results in this paper shown arithmetic modulo 113 (Power et al. 2022 Nanda et al.] 2023). This is a family of supervised learning tasks where two one-hot encoded inputs representing integers b < p used to predict the target y a mod p. where is some binary operation and p is a prime number. In most of our results, the binary operation is addition, but show additional results with multiplication and subtraction.
Modular arithmetic tasks are characterized by a binary operation and a dataset size, with different behaviors being observed for different dataset sizes on the same binary operation. In these settings, we describe the dataset sizes as the percentage of the 113 possible pairs that are used for training, with the rest of the data being used for testing as Nanda al. (2023) and Power al. (2022). Our main results use a 40%/60% train/test split but we also include results using 60%/40% and 70%/30%. The input integers are represented as one-hot vectors.
IL Sparse parity. We also validate some of our results on the Sparse Parity task outlined Barak et al. (2022. This is supervised learning setting where the target is the parity of k bits out of binary vector of length n, with it < n. In this work we use 2000 samples, split evenly between train and test data and we describe instances of this task by specifying the values of n and k.
III. MNIST. Finally, we provide some results on subset the classic image classification dataset MNIST (Deng) 2012). For our experiments, we use a subset of 200 training samples from the training set as in Liu et al. (2023b), with evaluation on the full test set.
2.2 MODELS
We study the grokking phenomenon on these datasets using 2-hidden layer multi-layer perceptron (MLP) of width 200 as in Liu et al. (2023a) and one-layer transformer with 4 attention heads as Nanda et al. 2023) and Power et al. (2022). We train both of these models in a full batch setting, using ReLU activations and cross-entropy loss with Adam W and SGD, well our own variants of these optimizers, 1. A and JSD. Unless specified otherwise we set the weight decay parameter À For modular arithmetic datasets, inputs are concatenated as the input of the MLP resulting in 226 dimensional vector, and treated as separate tokens in the case of the transformer.
3 SOFTMAX COLLAPSE: FLOATING POINT ERRORS PREVENT GROKKING
Given our current understanding of grokking, it is surprising that it happens without regularization for some dataset sizes, but regularization becomes crucial as dataset size decreases (Power et al., 2022). In this section we highlight that looking at datasets at the boundary of these two regimes reveals that without weight decay, grokking sometimes starts before abruptly stopping (Fig. 2 We show that this is caused by floating point errors in the Softmax that lead the gradients from large fraction of the samples to become zero. We refer to this phenomenon Softmax Collapse
3.1 SOFTMAX COLLAPSE
In modern neural network implementations, Floating Point (FP) arithmetic is ubiquitous for repre- senting and computing parameters, activations, and gradients. While FP numbers enable efficient decimal computations, they introduce numerical inaccuracies. This section focuses absorption errors, as a specific class of FP arithmetic failure. We will use the symbol to refer to equality under FP arithmetic.
Definition 1 (Absorption Errors). Let a, b E R {0} be floating point numbers in system with base 3 and p significand bits. Denote their exponents by en and eb, respectively. An absorption error occurs in the computation of a + b (denoted b = a) if
C, -s 2pp
In this case, after exponent alignment, the significand of is shifted right by at least digits, and b cannot be represented in the available precision, resulting in + =
Intuitively, absorption errors can occur during FP addition when operands have significantly differ- ent magnitudes. For float32 the base 3 is 2 and p 24 bits, meaning that adding any number smaller than 2 (x--11) 2-23 to 1 will leave unchanged. 2 2-23 is the machine epsilon for float32.
3

---

Published as a conference paper at ICLR 2025
100
Accuracy (9%)
Tal acc ibotls
60-
The arc -bbet37
That acc :bbttd
Thin aecurades
40
55% Soltrax
5K
16k
15k
20k
Epoch
Epach
Epach
(a) 40% training data
(b) 60% training data
(c) 70% training data
Figure 2: As dataset size increases (subplots a to c), MLPs trained on modular addition begin generalize without regularization until this is stopped by SC making the gradient from a large fraction of the samples equal to zero. This stopping point comes earlier for float32 than float64 and with small enough datasets it comes before the model makes any progress on test accuracy.
Absorption errors in the Softmax. The Softmax function is fundamental component in nu- merous deep learning architectures, serving as an activation function or a key element in attention mechanisms. In this case, we focus its application within the Softmax Cross-Entrope (SCE) loss: Definition 2 (Softmax Cross-Entropy (SCE) loss), For a neural network f and data point x with label y, we define z (x) and by as the logit corresponding to the true class y We express the SCE loss as well as its equivalent numerically more stable formulation as:
EScE( --log
²xyy
-y
+ max(2) log
Σ 2k -mxx(z)
Eke E
Unfortunately, even the rightmost (comparatively more stable) variant does not address this problem, since the kind of FP errors discussed in this work appear in the sum. While the Softmax function outputs are bounded between 0 and 1, the intermediate calculations involve summing exponentials of both positive and negative logits. These values can span several orders of magnitude, particularly in scenarios with large logits where the loss approaches zero. This wide range of values creates conditions that lead absorption errors leading the phenomenon we call Softmax Collapse.
Definition 3 (Softmax Collapse (SC)). A specific case of absorption error occurs when, for given sample x, the logit from the correct class Zy is significantly larger than the logits for all other classes. This floating point absorption of smaller terms, which call Softmax Collapse, occurs when:
Σ t = c
(2)
in which case the SCE loss becomes:
(scR(((x),)) log
(3)
Thus, during SC the loss becomes identical to zero. Furthermore, for the correct class, the gradients become zero as well:
aLSCE
E0c (c=y) = (c=w)
(4)
Oze
Ei_1 0th
While weights that contribute the wrong classes can still get negative updates, we show that
disappearance of the gradients from the correct classes is enough to inhibit grokking (Fig. (2). We
validate this in App. [B with an explicit intervention, showing that artificially setting the gradients from the correct class to stops generalization in very similar way to what observe in Fig
3.2 EVIDENCE OF SOFTMAX COLLAPSE IN GROKKING TASKS
Grokking is often studied using dataset sizes for which the delay in generalization is significant, which is usually when the dataset small but just large enough that generalization is possible. In
4

---

Published as a conference paper at ICLR 2025
this regime, regularization seems necessary for grokking and improvement in test performance observed without it (Nanda et al. 2023). However, a fact that has received less attention that grokking can happen without regularization if the dataset is large enough (Power et al. 2022).
Here we hypothesize that as the size of the dataset decreases, overfitting becomes easier and Softmax Collapse (SC) happens earlier. To quantify this, we train an MLP without regularization addition using different levels of FP precision, and calculate at every training epoch the fraction of samples that result in SC per Eq. (2). The results support our hypothesis that SC is responsible for the model's failure to generalize (Fig. 2 Specifically, we see that generalization stops when SC begins and that this happens earlier under float32 than under float64 (Fig. 2b). Furthermore, this point is reached earlier as the dataset size decreases until it is reached before making any progress in the test accuracy, resulting the common picture no grokking without regularization (Fig. 2a
modular
3.3 PREVENTING SOFTMAX COLLAPSE LEADS TO GROKKING
To validate the importance of FP errors in stopping grokking, show that methods to avoid SC lead to generalization on all the common grokking tasks on both MLPs and transformers. We introduce the following methods postpone the appearance of FP errors.
Increasing floating point precision. The simplest way to avoid SC is to extend the FP precision from float32 to float64 for the Softmax calculation. We see in Fig. 2 that networks trained using floa164 in the Softmax face SC later in training which allows for further increase in test perfor- mance. Conversely, using float 16 leads to SC earlier in training, leading to lower test performance. While this approach works as expected, FP precision cannot be extended indefinitely allow for generalization as seen in the lack of grokking in Fig, 2a
StableMax Cross Entropy (StCE) Loss. As demonstrated above, SC is caused by adding the exponentials of very large positive and negative logits in the Softmax. To avoid these extreme summands, we propose using a softer version of Softmax to transform logits into probabilities before calculating the CE Loss:
Definition 4 (StableMax). We introduce numerically stable version of the Softmax as:
y
y=
s(x,) [(1))
4
StableMax(z,)
(5)
2
where
1
x +
if 2 0, - if <0
s(x) =
(6)
0 Figure 3: 8(x) vs. e
As seen in Fig. 3 s(.) is a simple ramp function that scales linearly instead of exponentially when 0 and also approaches 0 more slowly than the exponential function when < 0. This is similar to the Softplus function (Dugas et al., (000) but approaches O more slowly with negative logits, further reducing the risk of absorption errors.
Proposition 1. StableMax is modified Softmax, i.e. StableMax (xx) moft (g (a:)) where
log(a + 1)
if= if <
g(x)
log( x + 1)
The proof of this Proposition is presented in App. A We then define the numerically stable analogue of LSCE LficE(f(x), v) Log(StableMax(- 2,)), where Zy again corresponds to the logit of the true class y.
To show that SICE indeed addresses the problems posed by SC, we repeat our experiments in Sec. 3.2 by replacing Softmax with StableMax. Our results, presented in Fig. indeed show that StableMax leads to grokking in commonly studied settings without regularization. Notably, this happens while the norm of the weights increases substantially (Fig. 4 middle). This suggests that while weight decay may lead to both grokking and decreasing weight norm, the decreasing
5

---

Published as a conference paper at ICLR 2025
100
aduk
100
40K
prodact mod
Spane partty norm
Accuracy (56)
Accuracy (96)
60-
thd acalton mod 1:3
Weight nomn
60-
20k
product
mos
Tisk ecCurity 2 FOS Test scturary rendom
20
(n-40) K-31
Thain soc.
20kc
40k
60k
50k
20k
40k
60k
80k
Epoch
Epoch
Epoch
Figure 4: (left) Grokking with SiCE loss and no regularization on three common grokking datasets using an MLP with 2 hidden layers of width 200. We use 40% of all pairs modulo 113 which is the same setting as Fig. [a where regular SCE gets stuck at random level performance (random level 50% for sparse parity). (middle) Evolution of model weight norms during training for the same models and tasks. This shows that grokking induced without weight decay does not follow the commonly observed trend of rapidly decreasing weight norm during generalization. (right) Changing input representations turns modular addition into regular machine learning tasks with train and test accuracy increasing in tandem, see Sec.
weight norm is not necessary for grokking. Overall, these results i) provide additional evidence for the importance of SC in preventing grokking, ii) suggest novel activation function to address this problem, and tii) show that regularization or weight norm modification is not necessary for grokking.
4 DIAGNOSING THE CAUSES OF SOFTMAX COLLAPSE
In the previous section we have shown that FP errors arise due to combination of low losses and large logits, and shown that when FP errors are mitigated, grokking can be observed in conditions where it previously was not. In this section, we dive deeper and ask why extremely low losses and large logits appear in the first place in grokking tasks. We identify two main causes for this tendency: (i) easiness of overfitting in grokking tasks, and (ii) training dynamic that gradients align with what we call naive loss minimization direction. After diagnosing the causes, the following section will these insights to develop an optimization algorithm that avoids NLM in the first place.
4.1 EASE OF OVERFITTING IN GROKKING TASKS
The first important characteristic of grokking tasks that lead SC is their ease of overfitting. It has been observed that as grokking datasets get larger, overfitting becomes harder, eventually leading regime where train and test performances increase tandem (Power et al., 2022 Nanda al.) 2023 Varma et al. 2023). It has also been shown that generalization can be delayed in the Sparse Parity lask by increasing the amount of noise in the input, which makes overfitting easier (Barak et al. 2022). Here we investigate the opposite effect: that by decreasing the dimensionality of the input the data becomes harder to memorize, removing the delay generalization.
To do this, we investigate the common grokking task of modular addition, but instead of the high- dimensional one-hot representations the input integers, we use a more compact binary. More specifically, we assign each integer a distinct random binary vector of dimension 14.
Results confirm our hypothesis, showing that as input representations are decreased in dimension, overfitting is prevented and models generalize without need for regularization (Fig. 4 right). This also shows that modular addition only induces grokking depending on the choice of representation. These findings highlight the importance of understanding the training dynamics beyond the point of overfitting (i.e. point of achieving 100% training accuracy), rather than focusing on the specifics of the modular arithmetic tasks as the key explaining the delay in generalization.
4.2
NAIVE LOSS MINIMIZATION
We next identify crucial training dynamic that commonly occurs in grokking tasks central cause for increasing logits and SC. We find that after reaching 100% training accuracy, gradient updates are dominated by an update direction term naive loss minimization (NLM). This direction

---

Published as a conference paper at ICLR 2025
Coptre Smilarty PILI
Costre Smitarty
enbessiW E.E
untersee V.
100%
100%
Epoch
Epoch
Epoch
(a) MLP without bias terms
(b) MLP with bias terms
(c) Transformer with bias terms
Figure 5: MLPs with (a) and without (b) bias terms trained on modular addition receive updates that significantly aligned with the direction of NLM beyond the point of overfitting. In (c) we show these results for a selection of parameters for our one layer transformer. We highlight the embed and unembed matrices well as the weights of the MLP. These are highlighted in the plot using the notation from Elhage et al. (2021).
does not change the model's decision boundary, but still decreases loss by simply scaling the logits of the predictions, in most cases through scaling of parameters (see below). This means that the logits will continue to increase until they inevitably lead to SC and zero terms in the training gradient. This stops the parameter updates in any direction, including NLM and any other useful component that would have been included in the overall gradient. We now define NLM formally, and proceed to discuss why might commonly be observed to deteriorate training in grokking tasks. Given the input E X, output y V predictor f parametrized by 0 E R that outputs logits (00x)) E R(>"), and a loss function L. we now define Naïve Loss Minimization.
Definition (Naive Loss Mininization (NLM)). A function dNLM : Rm - R" specifies direction of naive loss minimization if decreases the loss,
L( + (0):)) < ((f(O:)))
(8)
while satisfying for some c > 1:
f(O +
= c0) Yx E X,
where X denotes the imput space and
is the total loss over the training dataset.
We find that under a large class of models, namely those that demonstrate positive homogeneity, when training beyond 100% training accuracy the direction of the weights is an NLM direction.
Definition 6 (Positive Homogeneity (Lyu & Li 2020 A function f is positively homogeneous of degree L > 0 for all weights O inputs x, and scalars 0, it satisfies:
f(c0; x) c² f(0
(10)
When f a homogeneous neural network, corresponds to the number of layers.
In the case of homogeneous networks, training beyond 100% training accuracy, scaling the logits always leads to a decrease in the training loss. Therefore, dNLM(() = at for a > 0 is an NLM direction, as results in f(0 + (f): = f(11 + a);;z) (1 + f(0: x:), where the second equality follows from Eq. 1T))
Many neural network architectures, such as ReLU MLPs and transformers without bias terms, are positively homogeneous or appreximately homogeneous in the case of transformers (Merrill et al.) 2020 While more complex deep learning models with skip connections and bias terms are not homogencous, they have been shown to be quasi-homogeneouu (Kunin et al.) 2023) and in most cases - including all of the models in this work, the last layer is homogencous. This means that for non-homogeneous models scaling the weights of the last layer corresponds to direction of NLM.
The fact that the gradients converge to the direction of the weights has been studied in previous works (Ji & Telgarsky 2020: 2019 2018 Lyu & Li 2020 to prove that homogeneous networks converge in direction under gradient flow and gradient descent (GD), and they perform normalized margin maximization even beyond the point of 100% training accuracy (Lyu & I.i 2020). However,
7

---

Published as a conference paper at ICLR 2025
120
160
20
Accuraty
kotusi: Adamm
Terr act 5GD
1300
2000
3200
5000
Epoch
Epach
Log Spoch
(a) Transformer, subtract. mod 113
(b) MLP, addition mod 113 (c) Trade-off between L2 and SCE
Figure 6: Comparing LAdamW and LSGD with baseline optimizers and AdamW with weight decay on (a) a transformer trained on subtraction mod 113 and (b) an MLP trained on addition modulo 113. In (c) highlight the trade-off between L2 regularization and SCE loss, initially SCE loss is reduced at the cost of increasing the L2 loss but eventually the two losses decrease simultaneously (Sec. (Sec.(5.2) [..)).
we argue that gradient alignment also results scaling of the logits which can lead to SC and put an end to the margin maximization described in Lyu & Li (2020), when working with limited floating point precision. While we study delayed generalization. the link between training trajectories and generalization is already established in prior art (Birdal al. 2021 Andreeva et al. 2024)..
Evidence of naïve loss minimization. In practice, observe that in MLPs and transformers with and without bias terms, the gradients quickly become aligned with the direction of the weights after the point of overfitting (Fig- 5). Particularly for the later layers of the models, the cosine similarity between the parameter updates and the NLM direction goes up to 0.9 for the output layers. While models with bias terms are homogeneous and there no theoretical guarantee that scaling the weights will reduce the SCE loss, in practice we observe very similar behavior in MLPs with (Fig. (5b) and without (Fig. 5a bias terms. In the case of one-layer transformer. the alignment is stronger for the embed and unembed matrices but also substantial for the MLP weights (Fig. 5c).
5 MITIGATING NAÎVE LOSS MINIMIZATION LEADS TO GROKKING
While we have shown in Sec. 3 that avoiding numerical instabilities eventually leads to generaliza- tion, we can also target the NLM process that causes these numerical issues. To do this, we design an optimizer that only preserves the part of the gradient orthogonal to the direction of the weights.
5.1 .GGrdd AN OPTIMIZER TO PREVENT NLM
We propose a new optimizer, LGad (read "ortho-grad"), that updates the weights based only on the part of the gradient that is orthogonal to the current direction of the weights:
Definition 7 ( Grad). We propose the following update rule for a given iteration t E IN:
0 = 0. -ivi)/
(11)
where the orthogonal component of the gradient, C(O.). obtained by projection onto the hy- perplane orthogonal to the current weight vector:
0, VL(O.) 0.-
VLL(0)) =LC())))((())))
(12)
0. 0,
Proposition 2. Assuming V (L,( 0 33> 0 such that for any learning rate 0 < 1) < 3, laking the step v)V_j((C _((,, reduces the loss. In other words, any nonzero (iC(0,) is descent direction.
Sketch of the proof. We show that any [C(,) that (- ((),), vL())) < O. For a full proof we refer the reader to App.
{'((()) a descent direction by demonstrating
This projection of the gradient can be incorporated into different optimizers. Fig. 6a we show results for _AdamW and LSGD, the _Grad versions of AdamW and SGD respectively. These results show that .Grd optimizers lead to generalization without phase of initial overfitting, in contexts where improvement in test performance is usually observed without weight decay. We
8

---

Published as a conference paper at ICLR 2025
Piine pal Companers
Pttee pal Camporent
Tost LOSS
Jecap
SIID neight 900W
NLH diecciion
Principal Component
Component
1950
2000
(a) Training loss landscape
(b) Test loss landscape
Figure 7: Model trajectories in in parameter space projected 2D over the SCE loss landscape.
SGD with weight decay starts along the same trajectory as SGD decreasing the training loss (a) but increasing the test loss (b).
note that similar projections of the gradients have been used in other settings to mitigate the effects of momentum in invariant layers (Heo et al. 202T), stabilize training Wang al. (2024) one part in more complex optimizer (Kosson et al. 2024). We design .Grad as more precise intervention that directly prevents scaling along the NLM direction.
In Fig. 7 we compare the trajectories of models using SGD with and without weight decay to our new LSGD optimizer. SGD models start on similar trajectory. reducing the training loss but increasing the test loss, until the model with weight decay changes direction and starts minimizing both the train and test loss. In contrast, the model using LSGD moves directly direction that minimizes both the train and test loss. While SGD with weight decay eventually reaches a point of lower loss, note that LSGD reaches 100% test accuracy within 400 iterations (Fig. (6a). Beyond showing how LSGD prevents NLM, Fig. 7 also suggests that weight decay induces grokking by avoiding NLM. In the following, we highlight that the success of several methods to induce grokking can be explained from this perspective.
5.2 EXPLAINING THE SUCCESS OF EXISTING METHODS FOR GROKKING
In light of our findings, induce grokking. We find that these methods also lead to grokking by mitigating NLM and avoiding the FP errors that come with extremely low losses.
able explain the success of several previously proposed methods to
Weight decay. We have argued that the problem faced in grokking is that the of overfitting leads to NLM, which corresponds to scaling up the weights for homogeneous networks. Since weight decay corresponds to pulling back the weights along this same direction every step during training, it is unsurprising, given our findings, that it is the most reliable way to induce grokking.
To explain why generalization tends be delayed when using weight decay, opposed J. Gradd we look at it from the perspective of L2 regularization which is equivalent to weight decay for SGD. In Fig. 6c we see an initial phase where classification loss decreases, the cost the L2 loss. Eventually, the decrease in classification loss from NLM stops outweighing the increase in L2 loss, meaning that only updates that are not aligned with the NLM direction are followed. This explains why weight decay leads to generalization in grokking tasks but only after scaling along the NLM direction no longer decreases the overall loss. This balance between weight decay and classification loss is similar to the rotational equilibrium studied in Kosson et al. (2024).
We argue that the main roles of weight decay are preventing floating point errors and preventing NLM. This is in line with recent findings about the role of weight decay in deep learning (D' Angelo ctt al. 2023) which point to the fact that it increases the effective learning rate and avoids floating training in LLMs.
point issues when using
MSE loss on shallow networks. While cross-entropy loss can be reduced indefinitely by scaling the logits through NLM, this is not the case with MSE loss. When using MSE loss the logits can overshoot the target, meaning that larger logits often do not lead to lower MSE loss. This explains whylBarak
Barak et al. (022), Kumar et al. (2024). and[Lyu al. (2024) observed grokking with MSE loss without regularization. Interestingly. networks with more than one hidden layer do not generalize in these same settings [
Delaying generalization by scaling the weights While the lazy training dynamics described in Kumar al. (2024) explain important part of why scaling the weights delays generalization,

---

Published as a conference paper at ICLR 2025
show that the reason that regularization is often needed to exit this lazy training regime is that scaling the weights or the logits facilitates SC. In App. we show that the setting used in et al. (022b) to induce grokking on MNIST with SCE also induces SC which prevents further learning in the absence of weight decay.
RELATED WORK
Grokking. Power et al. (2022) introduced grokking and showed that weight decay can consistently induce it in algorithmic tasks. Nanda et al. (2023) were able reverse engineer the inner work- ings of a grokked transformer and found progress measures for grokking induced by weight decay. Chughtai et al. (2023) generalized the findings from Nanda et al. (2023) and showed grokked net- works use group representations to solve group composition tasks, although some of these findings were disputed in Stander et al. (2024) which propose that grokked networks learn a coset based algorithm for these same tasks. Mallinar et al. (2024) has shown that grokking is not specific to neural networks or gradicnt- based optimization and cannot be predicted from the training or test loss. Varma et al. (2023) argued that grokking is driven by weight decay favoring more efficient solutions and Liu et al. (2023b) hypothesized that the weight norm of the models needs to be in a "Goldilock's s zone" to generalize. Kumar et al. (2024) and al. (2024) connected grokking to transition between "lazy training" (Chizat et al. (0188) and feature learning, and Kumar et al. (2024) showed that this can happen without regularization in the case of shallow networks with MSE loss.
Grokking has also been described as phase transition by Zunkovië & Iliieski (2024). Lyu et al. (024) and Rubin et al. (2024). Humnayun et al. (2024) show that in many settings, neural networks undergo grokking- like transitions in their adversarial robustness. This aligns with the findings of Lyu & Li 2020) which attributed this increased robustness to a bias of SGD towards a max-margin solution which was proven for homogeneous models. Beck al. (2024) also connected grokking to the linear separability of the training data.
Numerical instability in deep learning. Numerical instability is a common issue in deep learning Kloberdanz et al. (2022). especially when dealing with mixed precision training D' Angelo et al. (2023). It is known that the Softmax function is particularly prone to numerical stability problems although this often comes in the form of overflow in the exponential (Kloberdanz et al.) 2022) and not from absorption errors in the sum as observed in this case. In the grokking setting, Nanda et al. (2023) showed that the slingshots observed Thilak et al. (2022) can be explained by very similar mechanism to the one involved in SC, although Nanda et al. (2023) do not use it to explain any grokking phenomena beyond these spikes that sometimes appear the training process in grokking tasks. We believe the slingshots observed in Thilak et al. (2022) could be a mechanism to prevent full SC, explaining why slingshots can lead to grokking without weight decay in some settings. This is further discussed in App. H Issues with numerical instability when training beyond overfitting with increasing learning rates were also observed in Lyu & Li (2020).
7 CONCLUSION AND DISCUSSION
In this work, we show that naive loss minimization (NLM) and floating point errors can explain why generalization is delayed in grokking and why it often does not happen without regularization. Using this insight, we are able to explain the success of existing methods to induce grokking. Motivated by our findings, we further design simple modification the Softmax that induces grokking by avoiding floating point errors and an optimizer that avoids the delay in generalization in grokking by preventing NLM.
Limitations & future work. While this work explains several surprising aspects of grokking set- tings, several questions remain. Notably, we focus our study of NLM on homogencous or approx- imately homogeneous models. A a formal characterization quasi-homogenous models could shed light on this kind of dynamics for models including skip connections and bias terms. Addition- ally, our explanation for why weight decay causes grokking could be enhanced by an analysis of its impact on the effective learning rate as a potential explanation for the sudden nature of grokking.
Acknowledgments. This work was supported by the UKRI Centre for Doctoral Training in Safe and Trusted AI [EP/S0233356/1] TB acknowledges support from the Engineering and Physical Sciences Re- search Council [grant EP/X01 1364/1]. TB was supported by UKRI Future I Leaders Fellowship [grant number MR/Y018818/11
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Rayna Andreeva, Benjamin Dupuis, Rik Sarkar, Tolga Birdal, and Umut Şimşekli. Topological generalization bounds for discrete-time stochastic optimization algorithms. Adv. Neural Inf. Process. Syst., 2024.
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit, Advances in Neural Information Process- ing Systems, 35, 2022.
Alon Beck, Noam Levi, and Yohai Bar-Sinai. Grokking the edge of linear separability. arXiv preprint arXiv:2410.04489, 2024.
Tolga Birdal, Aaron Lou, Leonidas Guibas, and Umut Simsekli. Intrinsic dimension, persistent homology and generalization in neural networks. Advances in Neural Information Processing Systems, 34, 2021.
Lénaïe Chizat, Edouard Oyallon, and F. Bach. On lazy training in differentiable programming. Advances in neural information processing systems, pp.
December 2018. ISSN 1049-5258.
Bilal Chughtai, Lawrence Chan, and Neel Nanda. A toy model of universality: Reverse engineering how networks learn group operations. In International Conference Machine Learning. PMLR, 2023.
Francesco D' Angelo, Maksym Andriushchenko, Aditya Varre, and Nicolas Flammarion. Why do we need weight decay in modern deep learning? arXiv preprint wrXiv:2310.044// 2023.
Haowen Deng. Tolga Birdal, and Slobodan lli.. Ppfnet: Global context aware local features for robust 3d point matching. In Proceedings of the IEEE conference computer vision and pattern recognition. 2018.
Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Process- ing Magazine, 29(6):141-142, 2012.
Jacob Devlin, Ming-Wes Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.). Pro- ceedings of the Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 4171 4186. Association for Computational Linguistics, June 2019.
Charles Dugas, Yoshua Bengio, François Bélisle, Claude Nadcau, and René Garcia. Incorporating second-order functional knowledge for better option pricing. In T. Leen, T. Dietterich, and V. Tresp (eds.), Advances in Neural Information Processing Syxtems, volume 13. MIT Press, 2000.
Nelson Elhage. Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatficld-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodci, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https:/taansfommeciiiuuu.pppppp wttp:/iiiiii.itau/o/aaaaaaeeeeee //fram/worr/inde....i..
Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.02677 2023.
Bycongho Heo, Sanghyuk Chun, Scong Joon Oh, Dongyoon Han, Sangdoo Yun. Gyuwan Kim, Youngjung Uh, and Jung-Woo Ha. Adamp: Slowing down the slowdown for momentum optimizers scale-invariant weights. International Conference on Learning Representations, 2021.
Ahmed Imtiaz Humayun, Randall Balestriero, and Richard Baraniuk. Deep networks always grok and here is why. arXiv preprint arXiv:2402.11553 2024.
Ziwei Ji and Matus Telgarsky. Risk and parameter convergence of logistic regression. arXiv preprint arXiv: 1803.07300 2018
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In 7th International Conference on Learning Representations, ICLR, 2019.
Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin Advances in Neural Information Processing Systems, volume 33, pp. 17176-17186. Curran Associates, Inc., 2020.
Eliska Kloberdanz, Kyle G Kloberdanz, and Wei Le. Deepstability: A study of unstable numerical methods and their solutions in deep learning. In Proceedings of the 44th International Conference Software Engineering, pp. 386-597, 2022.
Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay balances learning across neural networks, 2024.
11

---

Published as a conference paper at ICLR 2025
Alex Krizhevsky, Ily Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In Advances in neural information processing systems, volume 25, pp. 1097-1105, 2012.
Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the transition from lazy to rich training dynamics. In The Twelfth International Conference Learning Reprexentations, 2024.
Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetrie maximum margin bias of quasi-homogeneous neural networks. In The Eleventh International Conference Learning Representa- tions, 2023.
Ziming Liu,, ErieJ Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. In The Eleventh International Conference on Learning Representations, 2023a.
Ziming Liu, Ziqian Zhong, and Max Tegmark. Grokking as simplification: A nonlinear complexity perspective. In UniReps: the First Workshop Unifying Representations in Neural Models, 2023b.
Ang Lv, Ruobing Xie, Xingwu Sun, Zhanhui Kang, and Rui Yan. Language models" grok" to copy. arXiv preprint
2024.
Kaifeng Lyu and Jian Li. Gradient descent maximizes margin of homogeneous neural networks. In Inter- national Conference on Leariing Representations, 2020.
Kaifeng Lyu, Jikai Jin, Zhiyuan Li, Simon Shaolei Du, Jason D. Lee, and Wei Hu. Dichotomy of early and late phase implicit biases can provably induce grokking. In The Twelfth International Conference on Learning Representations, 2024.
Neil Mallinar, Daniel Beaglehole, Libin Zhu, Adityanarayanan Radhakrishnan, Parthe Pandit, and Mikhail Belkin. Emergence in non-neural models: grokking modular arithmetic via average gradient outer product. arXiv preprint arXiv:2407.20199, 2024.
William Merrill, Vivek Ramanujan, Yoav Goldberg, Roy Schwartz, and Noah A. Smith. Parameter norm growth during training of transformers. CoRR, abs/2000.00997, 2020.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learning Representations, 2023.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprini arXiv:2201.02177 2022.
Charles R Qi, Hao Su, Kaichun Mo, and Leonidas J Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 652-660, 2017.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Noa Rubin, Inbar Seroussi, and Zohar Ringel. Grokking as first order phase transition in two layer networks. In The Twelfth International Conference on Learning Representations, 2024.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(3):211-52, 2015. doi: 10. 1007/s 1263-015-0816.y.
Dashiell Stander. Qinan Yu, Honglu Fan, and Stella Biderman. Grokking group multiplication with cosets. In Forty-first International Conference on Machine Learning, 2024.
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Josh Susskind. The slingshot mech- anism: An empirical study of adaptive optimizers and grokking phenomenon. In NeuriPS Workshop, 2022.
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency, arXiv preprint arXiv:2309.02233 2023.
Mingze Wang, Zeping Min, and Lei Wu. Achieving margin maximization exponentially fast via progressive norm rescaling, arXiv preprint vrXiv22311.13387 2024.
Bojan Zunkovič and Enej llievski. Grokking phase transitions in learning local rules with gradient descent. Journal of Machine Learning Research, 155,,, 2024.
12

---

Published as a conference paper at ICLR 2025
APPENDIX
support main paper, App- A presents the proofs for propositions paper, App- includes additional findings that support our main results, and App. D provides further discussion on conditions that lead to grokking.
A PROOFS
Proof of Prop.
(x())
Soflmax (g (a:))
(13)
cg(xg)
log(s;+11
if 0,
2 +1)
1ng(-
(14)
if < 0
if 0,
Eg 2y+1
(15)
if <
Σn -
StableMax(z:).
(16)
Proof of Prop. To prove that any nonzero V Vi _((+) a descent direction, we need to show that
(L0,,, ((())) < O, assuming ViL(0.) # 0:
0 VL(O.)
VC(0.). VC
0, S 5..
(17)
0, 0,
Expanding this yields:
VV + VL(O.), 0.
0, VC
5..
(18)
0, 0.
Since the incquality is unaffected by the scaling of the left hand side, we can, without loss generality, assume that the gradients are normalized, Icading to:
VC(0), O. O. VL()))
31.
(19)
0. O,
Since 0: c((e).
8, denotes the projection of the gradient onto the space spanned the weights, (:-) will
measure the acute angle incidence and hence Eq. [9 holds, with equality ViL(0.) = O, which prevented by assumption. This proves that - V (().) is descent direction while being perpendicular to the weights.
We note that the Grad stops when VL where the gradient is in the same direction with the parameter vector. Vi((O.) 0 can also he the case if ((,))) = O. which corresponds the loss function being local optimum.
(l,) = O. If VC.))) * 0, this corresponds to the condition
B ADDITIONAL FINDINGS
B. FURTHER EVIDENCE THAT SC PREVENTS GROKKING
While SC leads the gradient from correctly predicted samples to be zero, it does not do this for the incorrect classes. To validate that setting the gradients from the correct classes to zero is enough to stop learning, we do this artificially for model that is generalizing and show that learning stops after this intervention. In Fig. we see that the baseline model shown in green generalizes. this is stopped epoch 6000 for model shown in blue, after we perform this intervention.
13

---

This problem has been Published as a conference paper at ICLR 2025
100
1821155 Ted
acouractes
40
50% Softmax
16k
2.k
30k
40K
50k
0
10k
20k
30%
40k
50k
10k
20k
20%
40k
50k
Epoch
Epoch
Epoch
(a) 40% training data
(b) 60% training data
(c) 70% training data
Figure 9: We show that the same dynamics observed in Fig. can be observed with a learning rate scheduler instead of AdamW. This shows that this is due an implicit bias of adaptive optimizers.
The intervention is implemented by multiplying the logits for the right classes by 0 at each step after epoch 6000.
100
Test acc.. float64 with artifially
80
lest acc. floated
acturacies
B.2 SGD WITH LEARNING RATE SCHEDULING
artificially induced
60
Collapse
To show that our results are not due the inductive bias of adaptive moments in optimizers like AdamW, we replicate some of the AdamW results using SGD with Icarning rate scheduler. Our scheduler lar to the one in Lyu & Lij (2020) except at each step the norm of the full observe
simi-
0
divide the learning
6k
gradient, instead of the loss. In Fig. that SC also puts an end grokking in this setting.
Epoch
Figure 8: Taking a model that would normally generalize (green) and artifi- cially inducing SC has a very similar effect to the one observed in Fig.
C EFFECTIVE LEARNING RATE
Unexplored in the main paper. NLM also has the effect of reduc- ing the effective learning rate. For a gradient update using regu- lar gradient descent 0, A: v/ [(O:) is casy see that [V.(e.)
200
|10++1 0
Aborbee pradiett PM
observed before when training beyond the point of overfitting, for example, (Lyu & Li (2020) addressed it by using a loss based learning rate scheduler to keep up with the gradient. Theoret- ically, alternative could be simply extend the duration of training. According to our hypothesis, training for long enough should eventually lead to gencralization on grokking tasks if we prevent SC. However, we find that another kind of floating point error can also appears in these settings, namely, gradient absorp- tion errors in the weights.
woght
200004
420038
600000
8E0000
Epecro
Figure 10: Gradient absorption errors during training addition modulo 113.
a weight w, gradient absorption emrrs happen when gradient update small enough that it leaves weight unchanged. Using the notation outlined in this paper this can be formalized as w - 8C at w. In Fig. [10 we show that this happens for an MLP trained with SGD modular addition using 30% of the training data. As the norm of the gradient decreases, the percentage of the gradients that are absorbed by the weights increases substantially. Note that the number of gradients that are exactly zero remains stable while the number absorbed gradients increases substantially.
This issue naturally mitigated by second order moments for adaptive optimizers like Adam and AdamW which is why they do not frequently appear. However, they do prevent from showing grokking with vanilla gradient descent without any learning rate scheduling.
C. ADDITIONAL WAYS TO INDUCE GROKKING
Beyond the interventions deseribed in the main text, highlight two additional ways induce grokking that validate hypothesis.
14

---

Published as a conference paper at ICLR 2025
6.0
Peining lovs
Teiring loss lest losa
Taiiing lows
Test losa
Test liss
5.5
2
4.5
Loos
Loos
4.0
2
3.0
1
2.5
1k
Epoch
Epach
Epoch
(a) StableMax
(b) Logit regularization
(c) Taylor Softmax
Figure 11: Train and test losses during grokking induced by three different interventions.
Logit norm regularization. Since we argue that uncontrolled scaling of the logits is responsible for delaying grokking and leading to SC, we validate that preventing this scaling of the logits by adding the norm of the logits the loss, leads grokking without additional regularization (Fig.
11b
Taylor approximation of the Softmax. We have introduced StableMax as a change to the Softmax that leads to grokking without regularization. The motivation behind this is to prevent values in the sum of the Softmax that are very large very close zero. To this end, replacing the exponential with any function that sub-exponential beyond a certain point should have a similar effect. To demonstrate, we perform a further ex- periment using the second order Taylor approximation of the exponential
Norm of Fourier Component
50k
40k
30k
10k-
e2
(20)
replacing the exp in the Softmax. Since the Taylor approx- imation decreasing for logit to avoid this part the function. We deem this version Taylor - Softmax. In Fig. [] we see results similar to the ones in Sec. 333 but showing the losses instead of the accu- racies well as results for two additional methods to induce grokking. Note that our implementation of Taylor Softmax (Fig. [lcc] [Icg introduces an additional implicit regularization simi- lar to the one in Fig. Iib due to the gradient flowing through the subtraction of the mean. While this effectively combines the effects of Fig. Taa and Fig. 11b leading to grokking faster than the other two methods, our main paper shows results using Sta as cleaner intervention that does not introduce this additional regular- ization effect.
0
10
20
30 40
50
subtract the minimum
Frequency
Figure 12: Fourier components of the weights of the output layer of an MLP trained addition mod 113. Grokking is induced via StableMax and without weight decay.
C.2 SOLUTION LEARNED DURING GROKKING WITHOUT WEIGHT DECAY
Weight decay has heen identified as potentially responsible for inducing the periodic structures in the weights studied in |anda et al.] (202)). In Fig. T2] we show that ML.Ps that grok without weight decay on modular addition show similar sparsity in Fourier space as the one observed in Nanda et al. (2023). While these are very superficial results, they suggest that these structures can emerge without weight decay-induced "clean up" phase described Nanda et al. (2023).
D FURTHER DISCUSSION ON CONDITIONS THAT LEAD TO GROKKING
D. LI REGULARIZATION AND GROKKING
While it has been observed that LI regularization can lead to grokking in some Settings,,Nandee et al. 22023 con- sistently found no grokking with L.1 regularization and transformers and this setting has received substantially less attention than weight decay.
We observe that NLM scales the weights along their current direction. This means that larger weights are scaled more than small weights. However. while the sign of the gradient from LI regularization depends on the sign of the weights, the magnitude of this gradient does not depend on the magnitude of the weights. This means that, particularly on deep networks or transformers with with large weights, L.1 can sometimes be insufficient prevent NLM and the subsequent SC.
15

---

Published as a conference paper at ICLR 2025
lalt S. cad
pa11
2013
10000
200.0
30000
18000
10:00
12000
20:00
20:00
10:00
20.00
3000
10:00
(a) MSE: hidden layer
(h) MSE: 2 hidden layers
CE hidden layers
Figure 13: The a parameter controls generalization in settings where it happens by default. This is the case for shallow networks with MSE loss as shown in subplot (a). However, in deeper networks (b) or networks with CE loss and no regularization (c), can control the time of over-fitting, but no value of enough to trigger grokking.
100
100
80
80
Accuracy (%)
Accuracy (%)
60
60
Training accuracy
Training accuracy
20
Test accuracy
20
Test accuracy
Softmax Collapse
Softmax Collapse
2000
4000
6000
8000
2000
4000
6000
8000
Epoch
Epoch
(a) MLP without weight decay
(b) MLP with weight decay.
Figure 14: Replicating the grokking on MNIST for weight decay setting from Liu et (2023b). We find that MLPs with weights scaled up by 100 operate at the "edge of numerical stability" and the absence of weight decay, SC eventually reaches 100%, preventing any further generalization. When using weight decay, the weight norm is reduced, mitigating SC and eventually allowing for further generalization as the SC rate drops from 100%.
D.2 DELAYING GENERALIZATION BY SCALING THE WEIGHITS
Scaling the logits can delay generalization but induce it. Liu et 2023a), Kumar al. (024) and Lyu (2024) showed that an a parameter multiplying the logits can increase or reduce the delay generaliza- tion. We highlight in Fig. [3 that this is true for cases where generalization happens even without changing the scale of the logits (o 1). However, in most cases when using deeper networks cross-entropy loss, models unable induce grokking for any value of
do not generalize default without regularization and
argue Sec. 5.2 that the observation in Liu al. (2023a), Kumar al. (2024) and Lyu al. (224)) of grokking without regularization are due to the inductive bias of MSE loss which prevents NLM and leads to grokking in some settings for shallow networks.
Grokking on MNIST. We replicate the setting from Liu al. (023b) of grokking on MNIST with cross- entropy loss and show that without weight decay, the scaling factor of the weights leads to significant FP errors, preventing grokking from happening until this is alleviated by weight decay.
While SC explains why weight decay is needed to get the jump in performance observed in Fig. 14b It could also explain why inducing grokking by scaling the weights less effective when using SCE. While when using MSE loss, Liu et al. (2023a) are able to induce full grokking from random level predictions to close to full training accuracy, the same does not seem to be possible when using SCE. In fact, we see in Fig. [4b that since the beginning of training the rate of SC approaches 100%. This could explain why the observations with cross-entropy loss not the ones predicted by the lazy training theories outlined in Kumar et al. (024 which do not take limited floating point precision into account.
16

---

Published as a conference paper at ICLR 2025
10
80
Accuracy
Train acc.
40-
acc.
Train LGrad
Test acc, wd+4
Test acc. Grad
Test acc. - w
Trann acc. wd-6
acc.
acc.
230
400
602
800
1000
6.0
600
600
2003
Epach
Epoch
(a) Sweep over values of weight decay
(b) .GGrad vs best performing wd model
Increasing weight decay (WD) for MLP trained modular addition with Adam reduces the delay in generalization up to point where WD prevents convergence Fig. 15a Without any tunable hyper-parameters and without WD, -Grad leads to grokking faster than the best model with WD Fig. 15b
Figure 15:
E LGrad AND WEIGHT DECAY
we provide a more in depth comparison increasing the weight decay multiplier leads to smaller delay in generalization, hut only up a point. In this concrete setting, weight decay multiplier of 8, prevents the model from fully generalizing (Fig. [Tap. We then compare the best value of weight decay in this setting LGrad, which does require any hyper-parameter tuning. Fig. [Sb] shows that 1 Grad leads faster grokking even when compared tuned value weight decay. Note that the models with weight decay overfit immediately before grokking while Grad reaches 100% train and test accuracies almost at the same time.
Fig. 15]
Grad and weight decay. Fig. 15a
highlights that
ALTERNATIVES TO StableMax IN PREVENTING SC
While any intervention that prevents SC should lead to grokking or generalization, Fig. [6 shows that scaling the temperature the Softmax not enough prevent SC
label smoothing does prevent SC and lead to some generalization, but at the cost of introducing another in-
Accuracy (%)
Tast arc. Label smoothing
ductive bias that prevents full generalization and leads qualitatively different behavior. By comparison, the sim- ple change introduced StableMax prevents leads to grokking, serving as a validation for our hypothe- sis that gradient descent Icads to grokking by default, un- less this stopped by SC.
Temperatur scal
Test arc. Stabiomax
20
10002 20000
30003
43200
50000
Epochs
G StableMax AND LGrad IN REALISTIC SETTINGS
Figure 16: StableMax prevents SC and leads to grokking while temperature scaling with T 1e5 only gradually delays SC, and label smoothing does prevent SC but at the cost of keeping the model from fully generalizing.
While StableMax and L LGrad are designed as interven- tions to show that preventing SC leads to grokking and preventing NLM leads generalization (Fig. D. in this section explore if these methods applicable in GPT2-small or ResNets trained on image classification. We
more realistic settings like language modeling train GPT2-Small for 1 epoch on WikiText- 103 using a batch size of 16, a block size of 512, a learning rate of 5e - 4 and a weight decay of 0.01 using AdamW. W. The architecture is the regular GPT2- Small architecture from Radford et al. 2019), trained with cosine schedule and 1000 steps of warm-up.
2015), baseline is ResNet18 with SCE loss trained with SGD 0.9 momentum and le weight decay. We use standard data transformations random crop and random horizontal tlip and a step learning rate scheduler every 30 epochs for a full training run of 100 epochs. With respect to this baseline we report results replacing the Softmax with StableMax in the loss function, well as replacing SGD with SGD. Since test labels for Imagenet- k are not publicly available, we the validation set set and tune parererarameters fraction of the training set.
For CIFAR10 CIFARI00 and Imagenet- k (Russakovsky
as
17

---

Published as a conference paper at ICLR 2025
Aetiemac sotrac borel
Grad
Adtreman aluntan
Accummyy P0
alto Setrra
Vel L.red
Train ace Slab emac
Trin acr softmax
200
120
Lokk 2890
Step
Epoch
Epoch
(a) GPT2-Small on WikiTex12
(b) ResNet18 on CIFAR100
(c) ResNet18 on CIFARIO
Figure 17: Comparing Stablemax and _Grad AdamW with SCE on text data Fig. 17a and image data Fig. [7c] [ic For the GPT2-small results in Fig. 17a
we also include the results of replacing the
Softmax in the attention mechanism with StableMax.
Method
CIFAR10
CIFARI00
ImageNet-1k
WikiText-103 (Top-5)
Softmax CE
87.17% + 0.2
59.98% + 0.4
69.33% + 0.04
60.48% + 0.04
Stablemax CE
% + 0.2
60.63% + 0.4
65.87% + 0.22
51.85% + 0.47
LGrad
87.22% + 0.2
62.69% + 0.1
68.95% = 0.03
59.64% = 0.04
Stablemax Attention
58.52% + 0.04
Table I:
For the methods introduced in this paper, we report accuracies with standard deviations across five seeds for the CIFAR datasets and three seeds for Imagenet- lk and WikiText-103, We report Top-5 accuracy in the case of WikiText-108
H AND THE SLINGSHOT EFFECT
Thilak al. 2022 observed that spikes in the training loss appear when training on grokking tasks with
adaptive optimizers like Adam, and that these spikes can lead to generalization without weight decay. Although Nanda et al. (2023) showed that slingshots are not necessary for grokking, it is still unclear what mechanism of adaptive gradient optimizers induces this behavior and why it leads to generalization. In light of the results in this paper, we believe that slingshots could lead to generalization because they prevent full SC. Nanda et al.) (2023) pointed out that something like SC could be responsible for these slingshots. One possible mechamism would be that zero gradients for some samples due to SC rapidly diminish the second-order moments Icading to large update or slingshot which moves the model away from full SC, although more research would be needed to properly show this.
While related to work, slingshots timizers and can allow grokking. In contrast, we identify SC as very specific issue the Softmax that can affect any model trained with SCE, not only the ones trained with adaptive optimizers. Additionally SC pre- vents grokking whereas slingshots can lead to it. Wether and how slingshots are cause by SC remains an open research question, with some supporting evidence from Nanda et al. (2023) which show that slingshots can disappear when using float61.
different kind of instability which only appears with adaptive op-
I ADDITIONAL DETAILS ABOUT FLOATING POINTS
Beyond our main results, we found that some cases, grokking could stopped before SC due the parameter in Adam being too large. While the e term is designed to give numerical stability to the gradients, in settings with extremely low losses and gradients, the second order moments can be dominated by the e term, putting an end to learning where it would have continued with smaller value. This echoes the results in Thilak et al. (2022) which shows that increasing halts slingshots and grokking, with Nanda al. (2023) also alluding to the parameter being important some cases.
Surprisingly, we also found that simple re- implementation of torch.nn. f functional.log fomaa that does not use the official CUDA kernels can lead the models to keep learning beyond the point where the loss is exactly 0 and some gradients should be with appropriate calculation. outperforming the official implementa- tion for grokking tasks. Learning eventually also stops in this setting and this seems more like quirk of how gradients are calculated PyTorch in the absence explicitly defined backward pass.
18


Reviews:
[{"review_id": "DN8QHQH22q", "reviewer": "Reviewer_Mcra", "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.", "questions": "Refer to Weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors make a novel discovery regarding the crucial role of numerical stability in grokking and identify the Softmax Collapse (SC) phenomenon.\n- They reveal that SC results from naive loss minimization (NLM) of cross-entropy loss.\n- 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).** \n  - 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?\n  - In transformers, softmax is applied in all self-attention layers, not just the last layer. Does SC also occur in these layers?\n\n- **StableMax.**\n  - 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.\n\n- **Naive Loss Minimization (NLM).**\n  - 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$.\n  - It seems that this explanation holds only for cross-entropy loss. For squared loss, NLM does not appear to hold.\n\n- **$\\perp$Grad.** \n  - 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.\n  - 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?\n\n\nReferences:\n\n[1] Ji and Telgarsky. Risk and parameter convergence of logistic regression. COLT 2019.\n\n[2] Ji and Telgarsky. Gradient descent aligns the layers of deep linear networks. ICLR 2019.\n\n[3] Ji and Telgarsky. Directional convergence and alignment in deep learning. NeurIPS 2020.\n\n[4] Wang et al. Achieving Margin Maximization Exponentially Fast via Progressive Norm Rescaling. ICML 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ndhflo7nus", "reviewer": "Reviewer_sm8G", "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.", "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). \n\n\n2. Could the authors clarify why the proposed adjustments should be prioritized over traditional methods like weight decay?\n\n\n3. Can the authors elaborate on how the proposed optimizer relates to [2]?\n\n\n4. 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?\n\n\n5. For Figure 2, could the authors specify the task on which the models are trained?\n\n\n6. In Figure 6, do the authors have experiments demonstrating the impact of different weight decay values on the delay in generalization?\n\n[1] Lyu, Kaifeng, and Jian Li. \"Gradient descent maximizes the margin of homogeneous neural networks.\" _arXiv preprint arXiv:1906.05890_ (2019).\n\n\n[2] Heo, Byeongho, et al. \"Adamp: Slowing down the slowdown for momentum optimizers on scale-invariant weights.\" _arXiv preprint arXiv:2006.08217_ (2020).\n\n\n[3] Kosson, Atli, Bettina Messmer, and Martin Jaggi. \"Rotational equilibrium: How weight decay balances learning across neural networks.\" _arXiv preprint arXiv:2305.17212_ (2023).", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper is well-written and easy to follow.\n- The research problem is conceptually well-motivated and aligns with current challenges in understanding grokking.\n- 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.\n\n2. **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.\n\n3. **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.\n\n4. **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6ZPtJa18X8", "reviewer": "Reviewer_yRJ9", "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.", "questions": "- What’s the difference between absorption error and Roundoff error?\n\n- Why setting the gradient of correctly predicted (overfitted) samples to zero prevents generalization? Is there an intuitive explanation?\n\n- In Figure 1c, can grokking still occur after careful hyper parameters(initialization scale, learning rate) tunning?\n\n- [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?\n\n[1] Progress measures for grokking via mechanistic interpretability\n\n[2] The Slingshot Mechanism: An Empirical Study of Adaptive Optimizers and the Grokking Phenomenon", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "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?\n- 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.\n- Sec 4.1 compares one-hot encoding with binary and scalar representation to show overfitting is harder, which is not fair: \n    - 1. the one-hot vector does not have any information of the integers, but binary and scalar vectors do contain some info. \n    - 2. the model trained on binary and scalar vectors has fewer parameters, making overfitting harder. \nA 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eojJrRxiUn", "reviewer": "Reviewer_NUmH", "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.", "questions": "- How do other factors such as initialization scale, learning rate, etc. interact with grokking?\n- Do the results shown in Fig. 4 correspond to one the settings shown in Fig. 2?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nThe 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.\n\n- The phase ‘grokking tasks’ can be reworded as ‘tasks where grokking is commonly observed’, since there is no formal definition of grokking tasks.\n- The readability can be improved by better placement of the figures (same page as the corresponding discussion).\n- In lines 302-304, the digits 1- can be replaces with i) for better readability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CR2wP18FTi", "reviewer": "Reviewer_2L9F", "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.", "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.\n\n2. Is ⊥Grad beneficial for tasks where overfitting is not prominent? Understanding if this method generalizes to other contexts could add depth to the findings.\n\n3. ⊥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? \nThe paper would benefit from additional experiments in varied settings to benchmark ⊥Grad against standard optimizers and even on non grokking tasks. \n\n4. 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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\n**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.\n\n**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.\nThe practical necessity and advantage of StableMax might be limited in many standard applications.\n\n**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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)