Egalitarian Gradient Descent: A Simple Approach to Accelerated Grokking
Authors: Ali Saheb Pasand
Elvis Dohmatob
Venue: ICLR 2026 Poster (2026)
Keywords: grokking
optimization
generalization
acceleration
Primary area: optimization
TL;DR: We show grokking arises from asymmetric speeds of gradient descent along different principal directions of the gradients, and we propose egalitarian gradient descent which significantly accelerates grokking by equalizing speed along directions.

Abstract: Grokking is the phenomenon whereby, unlike the training performance which peaks very early on during training, the test/generalization performance of a model stagnates over arbitrarily many epochs and then suddenly jumps to usually close to perfect levels. In practice, it is desirable to reduce the length of such plateaus, that is to make the learning process "grok" faster. In this work, we provide new insights into grokking. First, we show both empirically and theoretically that grokking can be induced by asymmetric speeds of (stochastic) gradient descent, along different principal (i.e singular directions) of the gradients. We then propose a simple modification that normalizes the gradients so that dynamics along all the principal directions evolves at exactly the same speed. Then, we establish that this modified method, which we call egalitarian gradient descent (EGD) and can be seen as a carefully modified form of natural gradient descent, groks much faster. In fact, in some cases the stagnation is completely removed. Finally, we empirically show that on classical arithmetic problems like modular addition and sparse parity problem which this stagnation has been widely observed and intensively studied, that our proposed method removes the plateaus.

Published as a conference paper at ICLR 2026
EGALITARIAN GRADIENT DESCENT:
A SIMPLE APPROACH TO ACCELERATED GROKKING
Ali Saheb Pasand
Elvis Dohmatob
McGill University & Mila Institute
Concordia University & Mila Institute elvis www.attooecooncc...ooooooooooooe ca
al i. sahebpasande mcg1 1
ABSTRACT
Grokking is the phenomenon whereby, unlike the training performance which peaks very early on during training, the test/generalization performance of model stagnates over arbitrarily many epochs and then suddenly jumps to usually close perfect levels. In practice, it is desirable to reduce the length of such plateaus, that is to make the learning process "grok" faster. In this work, we provide new insights into grokking. First, we show both empirically and theoretically that grokking can be induced by asymmetric speeds of (stochastic) gradient descent, along different principal (i.e singular directions) of the gradients. We then propose a simple modification that normalizes the gradients so that dynamics along all the principal directions evolves at exactly the same speed. Then, we establish that this modified method, which we call egalitarian gradient descent (EGD) and can be seen as a carefully modified form of natural gradient descent, groks much faster. fact, in some cases the stagnation is completely removed. Finally, empirically show that on classical arithmetic problems like modular addition and sparse parity problem which this stagnation has been widely observed and intensively studied, that our proposed method removes the plateaus
INTRODUCTION
Neural networks sometimes exhibit a striking training dynamic known grokking: after rapidly driving the training error to (near) zero, test performance can linger near chance for a long period before rising abruptly near-perfect generalization--oten without any change the optimizer or explicit early stopping (Power er al., 2022). Grokking has been observed across architectures and tasks, from algorithmic problems such as modular arithmetic (Gromov) 2023 Doshi et al., 2024) and formal languages to more naturalistic settings (Liu et al., [023] Nanda et al. 2023). Yet despite rapidly growing body of empirical and mechanistic case studies, we still lack principled account of why the delay occurs, what structural features crystallize at the transition, and how to predict or control it (see Section for a detailed review of the existing literature).
In this work, we examine grokking through the lens of the dynamics of eigen-spectra of gradients during optimization, and propose a simple modification of (stochastic) gradient descent which provably reduces the length of the platcau without compromising the level of generalization of the model at the end of training.
Contributions. Our main contributions can be summarized as follows.
Egalitarian Gradient Descent (EGD). We propose novel and simple method for accelerated grokking and study its properties both theoretically and empirically. Our method operates by modifying the gradients at each layer so that the principal directions (aka singular directions) are conserved, but the speed of evolution of the dynamics along each of these directions is the same. This stabilizes the training by reducing the effect of ill-conditioned loss landscapes (where gradients vary significantly in magnitude across different principal directions), leading to accelerated grokking.
An Effective Theory. We develop theory which shows that our proposed EGD method is guaranteed to drastically grok, compared to vanilla methods such as (stochastic) gradient
Code available /ttttt:///giiiii com/ http://www....aaaaaee////gaaaaa/ -tt.

---

Published as a conference paper at ICLR 2026
descent. We also make formal links to natural gradient descent. In fact, we show that proposed method (which can be seen as a carefully modified version of natural gradient descent) is simplified version of Grokfast (which operates by low-pass filtering of the gradients to boost weaker components). At a high-level, we also show that both methods have the same inductive bias of making the scales of the dynamics of the evolution of the parameters comparable along different important directions.
Empirical Verification. We extensive experiments toy problems (binary classification of anisotropic multivariate data) and arithmetic problems (sparse parity, modular addition, etc.) and show that our proposed method typically groks immediately the very beginning of the learning curve. To illustrate the computational efficiency, robustness, and generality of our EGD compared to benchmarks like Grokfast (Lee et al.) (024)), we also provide a vast array of experiments in the appendix on more realistic datasets (MNIST, CIFAR-10) and models (MLPs, CNNs, transformers).
Vanilia SGD
EGD (proposed method!
Column Nomalzztion
Addition Modulo 79
Addition Modulo 997
Addition Medulo 127
1.00
Accuracy
0.75
0.50
0.25
0.00
10
10a
10a
10a
10"
100
10a
10a
10a
10
100
10a
10a
10a
10
Epoch
Epach
Epoch
Figure 1: Results on Modular Addition for different values of the modulus p. Solid lines correspond to test accuracy and broken lines correspond to train accuracy. In all cases, our proposed EGD (egalitarian gradient descent) method groks after only few epochs, while vanilla (stochastic) gradient descent stagnates for long period before eventually grokking. We also include "Column Normalization", simplification of EGD which simply rescales the columns of gradient matrices by dividing by their L2 norm. Even this simplification seems to grok much faster than the baseline, vanilla (S)GD. Refer Section for details and Appendix for the hyper-parameters used.
Vanilla
(proposed method!
Cakumn
Multiplication Modulo
Multiplication
Multipiiationn Modulo p 127
1.00
0.75
0.50
0.25
0.00
10r
101
10a
10"
10c
10
10:
10'
10"
107
10:
10*
Epoch
Epoch
Epoch
Figure 2: Results Modular Multiplication for different values of the modulus Solid lines correspond to test accuracy and broken lines correspond to train accuracy. In all cases, our proposed EGD method groks after only few epochs, while all the other methods stagnate long period before eventually grokking. Refer to Section 3 for details and to Appendix for the hyperparameters used.
RELATED WORK
Grokking the phenomenon where models first overfit the training set and only much later un- dergo sharp jump in test accuracy- first documented by PPower al. (2022). Subsequent empirical studies broadened the scope beyond purely algorithmic datasets and analyzed conditions
under which grokking appears disappears, e.g., via loss foss-norm tradeoffs, optimizer choice and
dataset/regularizztion choices (Liu et al. 2023; 2023) Nanda et al. 2023 Notsawo et al. 2025 Thilak et al.
2022 Liu et al.) 2022 Davies al. 2023 Notsawo et al.
2023 Varma et al. 2023)
2

---

Published as a conference paper at ICLR 2026
Vanilla SGD
EGD fproposed ed method)
Column Normalization
Parity(nm400, km2) Retu 5M)
Parityinm 100, km3) Ret NNlwwdth=
Parityin=50, km4) ReW NNlwidth- 1000
Accuracy
0.6
0.4
100
10
10a
10a
10f
10a
10a
10a
10a
109
10:
10a
10a
Epoch
Epoch
Epoch
Figure 3: Results on Sparse Parity Problem. Solid lines correspond to test accuracy and broken lines correspond to train accuracy. All three plots show that our method (EGD) groks significantly faster than other methods. Refer to Section 5 for details on the experimental setup and to Appendix B for the hyperparameters used.
Histogram of eigenvelues of GG"
Histogram elgenvalues GGT Methed: Vanilla
Method: Varilla GD
Epoch 14999, Train acc: 1.0, Test acc: 1.0
Epoch 0, Tran acc: 0,01, Test acc: 0.01
120-
00-
30
20
0.5
0.2
0.4
0.0
0.2
0.4
0.31
1.0
Epenvatuve
x300
Elgenvatue
XI0
Figure 4: Ill-conditioned Gradient Spectra causes delayed generalization. We consider the problem of learning addition modulo 97 from data, with a two-layer ReLU neural network. At the start of optimization through the end, the gradient matrix G' for the hidden layer has poor condition number. Here, we see that the largest singular-vall value (corresponding to fast direction) is much larger than the smallest (corresponding vanilla (S)GD to stall for arbitrarily long times, leading to delayed generalization (see Figure I. Our proposed method, EGD (egalitarian gradient descent) forces all the singular values of to equal.
slow directions). This causes the overall dynamics of
Plateau Phenomena and Singular Learning Theory. Long before grokking (Power et al. 2022). prolonged training error plateaus were analyzed in information geometry and statistical-mehhnee treatments of neural nets. Plateaus arise near singular parameter regions-caused by permutation symmetries and redundancies -heee where the Fisher information degenerates and gradient flow becomes extremely slow (Wei et al., 20088; Amari et al. 2018).. Classical
studies of multilayer (soft-committee) networks likewise reported long transients and quantified their dependence on teacher-student alignment (Saad & Solla 1995). Singular learning theory gives general account via algebraic geometry, relating generalization behavior and marginal likelihood to model singularities (Watanabe 2009). More recently, ceaissiiallmeccaniie analyses have shown how input-data spectra modulate whether plateaus appear prominently at all (Yoshida & Okada 2019). While these plateaus need not coincide with delayed generalization
grokking, the underlying mechanisms degenerate Fisher spectra, symmetry breaking, and slow modes are highly relevant to grokking.
Grokking Modular Arithmetic. rich line of mechanistic interpretability work reverse eeverse-nngineers the circuits by which small transformers and MI.ss implement modular addition. Notably, Zhong et al. (2023) identify complementary algorithmic mechanisms ("clock" and "pizza") and circular embeddings that emerge at the grokking transition. Complementary analyses and constructive solutions for modular addition are provided Gromov] (2023). while Doshi et al. (2024) extend these insights modular polynomials (e.g., multi-term addition and related arithmetic), showing that trained networks converge toward these circuits near the onset generalization.
Grokking Kernel Escape. Another perspective grokking dynamical transition from lazy, kernel-like regime to a rich, feature-leande leauning regime. Kumar et al. (2024); Walker et al.) (2025) formalize conditions under which such transition yields delayed generalization without requiring
3

---

Published as a conference paper at ICLR 2026
explicit regularization. For modular addition, Mohamadi et al. (2024) give theoretical analysis explaining why early kerne-rregime learning cannot generalize under symmetry constraints, while later training escapes the kernel and finds small-norm, generalizing solutions. Aligned with this view, Varma et al. (2023) study grokking through the lens of circuit efficiency. This study shows that networks quickly learn memorizing solutions which generalize poorly, and the circuits that generalize well will take longer to form. When both circuits are formed, the generalizing circuits will be dominant in generating outputs.
Training Stability and the Optimization Lens. An orthogonal view emphasizes numerical stability and optimization dynamies as bottlenecks that can stall generalization; PPieto et al. (2025) argue that operating near the edge of numerical stability can induce grokking-like like delays and propose remedies that restore or accelerate test performance. Similar to this perspective, Thilak et al. (2022) identify a mechanism called "Slingshot" in which cyclic phase transitions in adaptive optimizers co-occur with grokking. The early spectral characteristics of learning curves have also been proposed as predictors of grokking, reducing the need for long training (Notsawo et al. 2023). In a more general study, Liu et al. (2022) develop an effective theory of representation learning, showing that generalization emerges specific zone for the weights of network called the "Goldilocks Zone". Grokking occurs when the weights enter this region, a narrow phase between memorization and confusion. In broader perspectives, grokking has been linked to phenomena such as double descent and emergent abilities (Huang et al. 2024; Davies et al. (002)). suggesting that delayed generalization can arise from competition between memorization and generalization dynamics. Recent thcoretical analyses have studied grokking in high- high-dimensional linear and logistic regression models under Gaussian assumptions, relating delayed generalization to spectral properties of the empirical covariance near interpolation thresholds (Beck et al. 2025 Levi et al. 2024). Taken together, these studies show that both data and representation dependent dynamics play roles in the emergence of grokking phenomena. As a result, grokking can be affected by architectural, optimization dynamics, and data-related interventions.
Grokking beyond arithmetic and Delay Mitigation. Grokking is not limited to algorithmic tasks. It has been observed in some computer vision and natural understanding tasks (Liu et al., 2023 Lee ctt al. 2024). Also, structural grokking has been shown to occur in language models, where models discover hierarchical sentence structures after extensive training (Murty et al. 2023 Zhu et al. 2024).
As more practical tasks exhibit grokking behavior, an important and interesting research question is how we can reduce the delay between memorization and generalization. Practical interventions have been shown be effective shorten remove the grokking delay (Lee et al. 2024 Lyle et al. 2025). Grokfas/ amplifies slow (low-frequency) gradient components via simple optimizer-side filters, consistently accelerating grokking across tasks and architectures (Lee al., (024). In Section
4.2 we discuss the algorithmic and conceptual benefits of our proposed method over this Grokfast.
Also, accelerating grokking has shown to be beneficial in a practical scenario in which data distribu- tions shift during training. To address this, Lyle et al. (2025) propose effective learning rate scaling and warming as a method to trigger and accelerate feature-learning dynamics during training which can both accelerate grokking and address the issue of primacy bias in continual learning. These dynamics-aware methods complement representation-side interventions and regularization/norm- control levers observed to modulate the phenomenon (Liu et al. 2023 Nanda et al. 2023).
WARM-UP: MOTIVATION FROM A SIMPLIFIED SETUP
We start with simple analytically solvable example where the structure of the data (relative strength of features, relative difficulty of training and test datasets) and optimization choices (learning rate, size of initialization) interact to provably produce grokking curves with arbitrarily long plateaus. The setting we consider here is directly motivated by the observations in Figure where delayed generalization in vanilla GS)GD is caused by ill-conditioned gradient matrices.
Data Distribution. Fix some small
0 and let
be centered Gaussian random
() 2 Let Ttrain E R22 be a folded version of obtained by
1 0)
vector with covariance matrix Σ
conditioning on the event |aTei| R Let Ftest be folded version of obtained by conditioning the event (=T v))=(1)) < 0. Here,
eil = | > 8, where e = (1, is the horizontal basis vector in
4

---

Published as a conference paper at ICLR 2026
E IR is a fixed unit- which makes angle E (0, 222) with e1. Thus, the condition number of the feature covariance matrix is 1/E > 1. This covariance ill-conditioning, captures, albeit in a simplified and instumeenal/conntollabee manner, what is going in Figure
Figure 5: A Toy Setup which Induces Stagnation in Gradient Descent (GD). Training data points correspond circles and test data points corre- spond to stars (middle region). The broken lines correspond the large margin of the training data (their separation is 2s), while the solid line is the ground-true truth recision-boundar x(1) = O. The variance of the slow feature xx(2) scales like & < 1. GD would quickly find linear model which per- fectly separates the training data but will take time of order 1/s find the ground-truth model which attains perfect test accuracy. See Figure 6
Feature x
0.2
-0.2
0.4
Fast feature x
The training data is (11,)1)...... (Min Vn. where y; == sign(x;;)) and the feature vectors at 's are iid copies of Ttrain- Here, 22(5) := x ej is the 1 component of the ith datapoint 24. For the test data, the feature vectors are iid copies of Ztest instead. The larger the value of s, the easier it is to quickly memorize the training data (perfect training accuracy). The situation is illustrated in Figure 5
Remark 1. In the above, there a shift between the distribution of the training data and the test data. In Appendix A.3 we also consider the setting where the training distribution coincides with the test distribution, the distribution of z constructed above.
Model. We consider linear model x + sign(IT where the weights vector á is obtained by minimizing the quadratic loss function
n w y1)² [=TT
1 , 1| -11 n
((w)
Σ 71
(1)
(1)
where X Raxd is the design matrix with Fi,-.., response vector. We choose this loss function because it leads to tractable analysis while retaining the same phenomenology we would get using the logistic loss function, for example.
........,.. Fa, and Y (y), (y1,-... -. ;,) E (+1)" is
3.1 VANILLA GRADIENT- DESCENT DYNAMICS
With step size y) the vanilla gradient-descent (GD) on the loss gives the following recursion
wk) wokk - 1) - x (Xw(k 1) - Y)/n Aw(k 1) + 6, with b = T/nn, A = E Σ
(2)
= = X
(3)
Here > 0 is sufficiently small stepsize learning rate, 1 is the iteration (aka epoch), and is the empirical covariance matrix. We can explicitly solve the above recursion to get (see the appendix)
w(k) = A (0) (I A )iotsss
(4)
where wots := X+Y - - is the ordinary least- squares solution. Thus, w(k) interpolates between the initialization w0) and the least-squares solution mater
Define constants = m, (s) (0, 1), m2 m²(s) > 1, a(s) E E. 3 = B(E) E (0. 1) by
mi := E x e1ll = (sss)Q((s) m2 := E[|z, e | 11 smi a = 7m2. 8 = nE, (5)
where so is the PDF of the standard Gaussian distribution, and Q == 1 - ① is its survival function. Let the initialization be w()) u (u,, U2), and define the following dynamical quantities
#4 == a & in + (1 a a mm / m2, v% = 3 u2 La == / / + =vz Th := uw/Vk.
(6)
The evolution of the test error is given analytically by the following result.
5

---

Published as a conference paper at ICLR 2026
Theorem 1. For large it holds w.h.p that: for any iteration k 1, (k)(k)) ~ arccoo arccos(r)), with r := p/7, Mal Lks
(7)
where p == cos 0, 7 7: vp² t p²), La == V EVZ.
(8)
Theorem 1 rrevelss that: the test error plateaus at value 100%, until k is sufficiently large that decreases below r, which point the error starts decreasing the rate arccos(r;) arccos(r).
Vanilla
100
Accuracy
75
50
-logz
25
10
10a
102
10a
10*
100
10a
102
10a
10a
Epach
Epoch
Figure 6: Grokking on the Toy Problem. Solid lines correspond to experimental results, while broken lines correspond
theory (Theorem [}. The initialization is w(0) is such that [[= Thus, the scalar C > 0 controls the size of the mitiillizatio.. Left. As predicted by Theorem 1 and Corollary [T] large initialization leads to delayed generalization in vanilla GD, i.e. long plateaus (of length k, x 1/E) of stagnation in the test performance, while small initialization attenuates (k: x log 1 I((re), where r x (). Note that for this problem, the ratio 1/e > represents condition number of the covariance matrix of the features. Right. Our proposed EGD (egalitarian gradient descent) method groks immediately: the generalization error jumps to perfect value after only a few iterations. Moreover, EGD appears to be completely insensitive to the scaling of the initialization, which is a desirable property real-world optimization.
We have the following important corollary.
Corollary 1 (The distribution-shiff case). Fix a constant >> 1. We have the following phase diagram: (A) Large Initialization. If |uz] is large in the sense that T == |2zl tan m22 /m > c, then the plateau length of the test error (i.e. the time to grok) is of order k. loR
(B) Small Initialization. If [uz] is small, then the plateau length is of order k. 1 log 1 TE
log
Corollary shows that a model trained via vanilla gradient descent will quickly converge to a decision- boundary which memorizes the training data (100% training accuracy, very poor test accuracy), but will take long time of order 1/s or log 1/s (depending on the size of initialization) before converges to the true increase to 100%. This result can be extended to vanilla stochastic gradient descent (SGD).
(Theorem [T, and only then does the test error abruptly
The Case of Same Train and Test Distribution. The setup (see Remark is analyzed in Appendix A.3 and we get the same order of magnitude, 1/(ne), for the length of the plateau in the case of vanilla GD.
The Issue with Vanilla Gradient Descent, One can show that for large sample size (n - oc), the empirical covariance matrix which modulates the dynamics
2 has the following approximation.
m2
0
Where m2 as defined in (5). Since m2 is fixed positive constant, the condition number of the RHS of order 1/s. This controls the rate at which the GD iterates w(k) converge to the least-squares solution via (4), which is in fact optimal (perfect test accuracy) here since we are in the infinite sample regime). Therefore,

---

Published as a conference paper at ICLR 2026
Insight #1. Grokking profile for this toy problem is therefore solely due to delayed convergence least- squares solution caused by ill-conditioned gradients.
3.2 ACCELERATED GROKKING VIA A MODIFIED GRADIENT SCHEME
Consider the following modified dynamics, which will later form the basis for our proposed method:
kk) wik 1) - (Xw(k 1) Y)/n,
with Σ = x x X-/n n as before. Expanding the above equation, the dynamics become
(kk) = w - - - X/n)i//nWwww - 1)) - x y/n) = aw is 1) + mpwole
with a == 17. What has happened is that the inverse FIM has killed the ill-conditioned Σ matrix which was multiplicatively damping the iterates. Solving the above recurrence gives
1 a (kk a² (0)) + Σ a' ymools = a² (0)) + muots = a w0(( + (1 - yWoos
(10)
The troublesome dependence on e has now completely disappeared from the picture. Moreover, we see that the convergence to Wots is now much faster than with the vanilla dynamics 4 namely
wk)) = A (0) [I- yWoot where == - E Thus, in (TO) an isotropic matrix al (spectral radius 17. fixed positive constant less than for any stepsize 1) (0, 1)) replaces
A= - 1 (1-0me 17m2 0
1-1E
which has spectral radius = 1- ne ~ 1, in 4 Therefore, under the modified gradient descent update rule (9). the iterates w=k) converge Wats exponential rate which independent of e. This leads to grokking after just few iterations, as seen in Figure
Insight #2. The modified GD update rule ( induces the desirable normalizing effect whereby the optimization dynamics has exactly the same speed along all principal directions.
4 PROPOSED METHOD: EGALITARIAN GRADIENT DESCENT
Motivated by the insights from Section 3.2 we now consider a general non-linear model (neural network, etc.) a general problem and let G be the gradient matrix for arbitrary layer. Thus, G has shape m x where m is the fan-out and p is the fan-in width for that layer. For example. if the layer is the hidden layer two-layer feedforward full-connected neural network, then m is the number of hidden neurons and p is the input-dimension. Consider the following transformation of the gradient matrix G:
(EGD)
F==F-1/2G (GGT)-1/24
(11)
where == GGT the empirical Fisher matrix and F-1/2 denotes square root. We call this transformation egalitarian gradient descent (EGD), that the above transformation leaves the left and right singular- vectors G unchanged but makes all the singular-valuee equal. Indeed, consider the singular-val valuee decomposition (SVD) of G:
that will become clear shortly. Observe
G = USV $1 Bannss
T + S2ugss +
(12)
where U = (m1)) and V = (vz); contain the left and right singular-vectors (aka principal directions) of G, and S = diag( 82,...) contains its singular-values. Then, (GGT) --12 U and so
G (GG (GGT)-1/20) - - UST = US-
= UVT
which has the same left and right singular-vectors as G but singular-values all equal to 1. This means that all through the optimization process, principal direction will evolve faster/slower than another. This justifies the name EGD (egalitarian gradient descent) given to (IT). and we shall show that drastically accelerates grokking.
7

---

Published as a conference paper at ICLR 2026
Remark 2. Note that in formula [11, we have implicitly assumed that G' has full-rank. In case of rank deficiency (which will happen if m > for example), we simply replace (GGT)-1 and S S by their Moore-Penrone seeudo-imersee.
Practical Considerations. The gradient G = UVT for our proposed EGD method is obtained by performing SVD the original gradient matrix G. This is the main computational cost incurred by our method. In practice, we turn off EGD and switch it for vanilla (S)GD once we detect grokking has occurred. by monitoring the validation loss. Moreover, additional experiments (refer to Appendices D) suggest that the SVD does not have to be precise, and approximate versions thereof, like randomized SVD (RSVD), work just fine, and lead to better stability and faster wall-clock time.
4.1 CONNECTION TO NATURAL GRADIENT DESCENT
The empirical Fisher information matrix (FIM) for any layer with gradient matrix is given by F = GG the same matrix which appears in (1I. is then easy that
| = (L/m) tr| FG]- (1/m) ((GGT)T1/22 GG G(GGT)
(1/m) tr
Thus, the Fisher-norm of the modified gradient matrix G is constant of motion of the dynamics induced by our proposed transformation (TT. Note however that our EGD proposed method is not
equivalent natural gradient descent (NGD) (Amari 1998 Pascanu & Bengio 2013 Ollivier et al. (GG
2017), which would correspond to G = F linked like so
G. Notwithstanding, the two methods
G =
F 1/2
G
(13)
EGD
FIM
NGD
Therefore, proposed EGD corresponds to whitened version of NGD. The effect of this whitening precisely to equalize the singular-values of the gradient matrix, leading to accelerated grokking.
4.2 COMPARISON WITH GRADIENT-FILTERRIN (GROKFAST)
The Grokfast method proposed by Lee al. (2024) for inducing fast grokking goes follows. Each row the gradient matrix G is replaced by g + F F(g), where F(g) is low-pass filtered version of g. computed by aggregating with large buffer of the past history of gradients. This has the desirable effect of boosting the frequency components of the gradient, and attenuating the high- geghffreyencc components thereof.
Now, let Cj := g Ju, be the jth component of g measured in the eigen-basi for GT G. From 12
(GGT)-1/2g (cI/si)u1 + (22/s2)u2 +...
(14)
This down-weight the components of g aligned with large s, the Biggh-frequecy"" directions so our EGD update inherits the filtering inductive bias of Grokfast (Lee al.) 2024) as a by-product. Crucially, unlike Grokfast, EGD equalizes the optimization speed across principal directions, yielding isotropic progress the eigenspace. EGD also enjoys the following important properties.
Memory. In contrast to Grokfast (Lee et al. 2024), which maintains a large buffer of past gradients, our proposed EGD method incurs no additional memory overhead beyond the current gradient (and, if used, running spectral estimate).
Simple and Hyperparameter-ffee. EGD introduces no extra tuning knobs. It therefore avoids the task-dependent, consecnnn hyperparameter sweeps that are important for Grokfast. The formula for EGD (II) lightweight, drop-in modification of (stochastic) gradient descent with a closed-form, per-step rescaling in the principal basis; it requires no schedulers, no momentum variants, and buffering.
Theoretical foundations. EGD comes with a spectral analysis guaranteeing equalized per- mode convergence rates and, consequently, accelerated exit from the test-error plateau (i.e., faster grokking). By comparison, Grokfast is a heuristic frequency filter without such guarantees.
8

---

Published as a conference paper at ICLR 2026
5
EXPERIMENTAL VALIDATION
5.1 SPARSE PARITY PROBLEM
This is a well-known hard problem in statistical learning theory (Barak et al. 2022). Also, recent works have shown that this problem induces grokking in (stochastic) gradient descent (Merrill et al. 2023). An instance Parily(n, k) of this problem is as follows. n. and k are positive integers with < n2 random k-element subset S of [n]] is drawn once and for all, and then N iid samples (1),,........ (IN, yN) are generated, where the are iid uniform r--it strings, and each Y/ corresponds to the XOR of the bits of restricted to the secret subset S, i. y ;= (-1)Enes "iy, The accuracy of a model f {0, 1)" + f-1,1} } is counting the proportion of points in a large held out test dataset (generated from the same distribution) have the true labels correctly predicted.
For the model, we consider two-layer ReLU network f(z) = sign(v' ((z))) trained by optimizing hinge loss and weight decay with using different optimization strategies, on the parity problem with different values of n and k: (n, k) E {4400, 2). (100, 3), (50, 4)). The batch size is set to 32. For reproducibility, information about low-level details like learning rate, amount of weight decay, etc. is provided in Appendix B We compare different optimization strategies: vanilla (stochastic) GD, our proposed method EGD [T] (applied only on the gradient matrix G for the hidden layer weights W), and an even simplified version of EGD where we replace each column of G.
The results are shown in Figure As predicted by our theory. EGD groks very early on in the optimization process (typically after only a few epochs). In contrast, vanilla (S)GD goes through an arbitrarily long plateau of stagnation of the test error before eventually grokking.
5.2
MODULAR ARITHMETIC
Another family of problems that exhibits grokking behavior is modular arithmetic. This class of problems has extensively been studied in vast array of papers
grokking (Power et al. 2022) Liu
et al.
2023 Lee et al. 2024 Mohamadi et al. [2024] Nanda et al. [2023] Notsawo et al. 2023 Zhong
al.
2023 Gromov 2023 Doshi et al. 2024: Prieto et al.; 2025). To formally define this class
of problems, an instance Mod(p, 0) defined by a prime modulus p and an operation o E {+, x over Zp. Training data consists of N i.i.d. samples (zt, 3t), which are fraction of all p² possible combinations. Here 2; = (ag, bx) is drawn uniformly from {0, 1..... p} x {0, I.... p), and the label calculated as yi = (a; by) mod p. We train two-layer ReLU network with cross-entropy loss and weight decay, using the same setup as in the sparse parity experiments. The complete set of hyperparameters is provided Appendix B Similar to sparse parity, we compare vanilla (stochastic) GD, our proposed method EGD (1]. and simplified column-wise variant of EGD.
As shown in Figures [ and 2 EGD groks considerably earlier than other methods, achieving high accuracy after only few epochs in both modular addition and multiplication tasks.
5.3 ADDITIONAL. EXPERIMENTS: COMPUTATIONAL EFFICIENCY AND ROBUSTNESS
Appendices C. D and present results of additional experiments in which we test the usage of EGD across different model types (MLPs, CNNs, transformers) and more realistic datasets (MNIST and CIFAR-10). We have also compared EGD with Grokfast (Lee et al. 2024) as the strongest baseline. These experiments continue to show that our method leads to f'aster grokking while being computationally efficient. Experiments in Appendix Elass show that EGD is robust in the presence of non-stationarity in the training data (encountered in reinforcement learning applications, for example).
6 CONCLUDING REMARKS
We studied grokking through the lens of gradient eigen-spectral dynamics and proposed Equalitarian Gradient Descent (EGD), a simple, hyperparameter-fin free modification of stochastic gradient descent that equalizes optimization speed across principal directions which can be efficiently estimated via inexact SVD (e.g., randomized SVD). By weightigg high-frequency components while preserving progress on slow, symmetry-alignat update that consistently shortens the acstraccuracy plateau without degrading final performance. Beyond empirical gains, our analysis clarifies how progress along low- now-frequency, task-aligned
modes, EGD provides principled, spectrum-aware

---

Published as a conference paper at ICLR 2026
directions governs the timing of the generalization jump, and it yields compact diagnostics that relate early gradient spectra to later test improvement. The method is optimizer agnostic, easy to integrate into existing training loops, and introduces no additional tuning knobs.
Our proposed EGD method is lightweight and straightforward to deploy, offering faster grokking with minimal computation, and no additional memory usage unlike Grokfast (Lee al. 2024).
Future Work. future, we also plan to study plug and play combinations with adaptive optimizers and weight decay, behavior under nonstationary data and curriculum schedules, and broader bench- marks beyond algorithmic tasks. On the theoretical front, beyond the simplified setup considered in Section [ it would be nice have a convergence analysis of EGD for complex models like MLPs and transformers. This is highly difficult, but we have some preliminary path to such an analysis using ideas from spiked random matrix phase transitions (Baik et al.) 2004).
REFERENCES
Kingma DP Ba J Adam al. method for stochastic optimization. arXiv preprint arXiv: 1412.6980 1412(6), 2014.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251-276. 1998.
Shun-ichi Amari, Tomoko Ozeki, Ryo Karakida, Yuki Yoshida, and Masato Okada. Dynamics of learning in mlp: Natural gradient and singularity revisited. Neural Computation, 30(1):1-33, 2018.
Jinho Baik, Gérard Ben Arous, and Sandrine Péché. Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices. Annals of Probability, 33, 2004.
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 Processing Systems, 35:21750-21764, 2022.
Alon Beck, Noam Itzhak Levi, and Yohai Bar-Sinai. Grokking at the edge of linear separability. In International Conference on Machine Learning, pp. 3307-3334. PMLR, 2025.
Roger Creus Castanyer, Johan Obando-Ceron, Lu Li, Pierre-Lue Bacon, Glen Berseth, Aaron Courville, and Pablo Samuel Castro. Stable gradients for stable learning at scale in deep reinforce- ment learning. arXiv preprint arXiv:2506. 15544, 2025.
Xander Davies, Lauro Langosco, and David Krueger. Unifying grokking and double descent. arXiv preprint arXiv:2303.061173 2023.
Darshil Doshi, Tianyu He, Aritra Das, and Andrey Gromov. Grokking modular polynomials. 2024.
URL httph:://spppi..../p///a/////ttt wwwpeiii.orr...ooooe.. wttp:///.gga//////////e// /406 03495
Andrey Gromov. Grokking modular arithmetic. 2023. URL
02679
Nathan Halko, Per-Gunnar Martinsson, and Joel A Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM review, 53(2): 217-288, 2011.
Yufei Huang, Shengding Hu, Xu Han, Zhiyuan Liu, and Maosong Sun. Unified view of grokking. double descent and emergent abilities: A perspective from circuits competition. arXiv preprint arXiv:2402.155 15175, 2024.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Tanishq Kumar, Blake Bordelon, Samuel Gershman, and Cengiz Pehlevan. Grokking as the transi- tion from lazy rich training dynamics. In International Conference Learning Representations (ICLR), 2024.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEFF, 86(1) 86(11):2273) (11())))278-2244 2002.
10

---

Published as a conference paper at ICLR 2026
Jaerin Lee, Bong Gyun Kang. Kihoon Kim, and Kyoung Mu Lee. Grokfast: Accelerated grokking by amplifying slow Jun 2024.
2024. URL httt:a//aarxdv.oorr/bbb/////2
v2,
Noam Itzhak Levi, Alon Beck, and Yohai Bar-Sinai. Grokking in linear estimators-a solvable model that groks without understanding. In The Twelfth International Conference on Learning Representations, 2024.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint 1rXiv:1908.03265 2019.
Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. Advances in Neural Information Processing Systems, 35, 2022.
Ziming Liu, Eric J. Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data. In International Conference on Learning Representations (ICLR), 2023. URL https:
org/ abs/ 2210. 01117 Spotlight.
llya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 10005:171 2017.
Clare Lyle, Gharda Sokar, Razvan Pascanu, and Andras Gyorgy. What can grokking teach us about learning under nonstationarity? arXiv preprint 2577v22507.20057, 2025.
William Merrill, Nikolaos Tsilivis, and Aman Shukla. A tale of two circuits: Grokking competition of sparse and dense subnetworks. arXiv preprint arXiv:2303. / 1873, 2023.
Mohamad Amin Mohamadi, Zhiyuan Li, Lei Wu, and Danica J. Sutherland. Why do you grok?
theoretical analysis on grokking modular addition. In International Conference on Machine Learning (ICML), 2024.
Shikhar Murty, Pratyusha Sharma, Jacob Andreas, and Christopher D Manning. Grokking of hierarchical structure in vanilla transformers. The 6lst Annual Meeting Of The Association For Computational Linguistics, 2023.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. 2023. URL http://aaaaai..oorab////23000 05217 v3, 19 Oct 2023.
Pascal Notsawo, Hattie Zhou, Mohammad Pezeshki, Irina Rish, Guillaume Dumas, et al. Predicting grokking long before it happens: A look into the loss landscape of models which grok. arXiv preprint arXiv:2306. 13253, 2023.
Pascal Notsawo, Guillaume Dumas, and Guillaume Rabusseau. Grokking beyond the euclidean norm of model parameters. In International Conference Machine Learning (ICML), 2025.
Yann Ollivier, Ludovic Amold, Anne Auger, and Nikolaus Hansen. Information geometric -enometrionngommtriic opti- mization algorithms: A unifying picture via invariance principles. Journal of Machine Learning Research, 18(18)11665 16,, 2017.
Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. arXiv preprint arXiv:1301.3584, 2013.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: General- ization beyond overfitting on small algorithmic datasets. 2022. URL
abs/ 2201. 02177
Lucas Pricto, Melih Barsbey, Pedro A. M. Mediano, and Tolga Birdal. Grokking at the edge of numerical stability. 2025. URL http:ttss//arxiv..rrt http://aaaaarrraeee/////////./..
David Saad and Sara Solla. On-line learning in soft committee machines. Physical Review E, 52 (4):4225, 1995.
11

---

Published as a conference paper at ICLR 2026
scikit-learn developers. kkbamm.ttteeetmmmt...mmmnnid..d Attp://sciklt-eear http://aabbeemccclss//eeeeeeeeee rateeeeeeeeeeeeeeeeeeeeeeaaaaeee wtw.re/ddd.eaaa.ttiiiiei. http:::ttttssttttt.......ataaaaa randemizedd
svd.html 2024. Accessed: 2025-11-18.
Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evci. The dormant neuron phe- nomenon in deep reinforcement learning. In International Conference on Machine Learning, pp. 32145-32168. PMLR. 2023.
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint arXiv:2206. 04817, 2022.
Tijmen Tieleman and Geoffrey Hinton. Divide the gradient by a running average of its recent magnitude. coursera: Neural networks for machine learning. In Technical report. University of Toronto, 2017.
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint arXiv:2309.02:55 02339, 2023.
Thomas Walker, Ahmed Imtiaz Humayun, Randall Balestriero, and Richard Baraniuk, Grokalign: Geometric characterisation and acceleration of grokking. arXiv preprint arxiv2550.112884 2025.
Sumio Watanabe. Algebraic geometry and statistical learning theory, volume 25. Cambridge university press, 2009.
Haikun Wei, Jun Zhang, Florent Cousseau, Tomoko Ozeki, and Shun-ichi Amari. Dynamics of learning near singularities in layered networks. Neural Comput.. 20(3), 2008
Yuki Yoshida and Masato Okada. Data-dependence of plateau phenomenon in learning with neural network statistical mechanical analysis. Advances in Neural Information Processing Systems, 32, 2019.
Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Xuekai Zhu, Yao Fu, Bowen Zhou, and Zhouhan Lin. Critical data size of language models from grokking perspective. arXiv preprint arXiv:2401.1044 2024.
12

---

Published as a conference paper at ICLR 2026
Appendix for "Egalitarian Gradient Descent: A Simple Approach to Accelerated Grokking"
CONTENTS
Proofs
13
A.I Analytic Theory for the Toy Problem
13
A.2 Important Lemmas
14
A.3
Theoretical Analysis of Toy Model: The Equi- distrribstiibiionnn Case
14
Hyper-parameters for Experiments in Section
17
Randomized SVD for Reducing Computational Cost
17
D
EGD in More Practical Scenarios and Empirical Comparison with Grokfast
20
EGD Helps Widely Used Optimizers Handle Non-Stationarity
23
Does EGD converge to different sup-space compared to vanilla SGD?
24
Implementation Details and Pseudo-codes
27
A PROOFS
A.1 ANALYTIC THEORY FOR THE TOY PROBLEM
Proof of Theorem I
First observe that the vanilla gradient descent dynamics can be expanded as
follows
w(k) w(k - 1) x (Nu(k - 1) V)/n = Aur(k 11) +
= Aru(k 2) + Ab + b.... 1²w((0) + (A^-1 +.... A A T)b
A4 + (1 A² AA) (0)(00) + (1- (1*t))n-13) -1
A0)(()) (I A A
Now, for N00 Σ) and 2g == 1-1 NO, La), we can write
v)) Etest (b (i)))) P((x (e)(TT))>> (k)) <0 (2)(=Te)( 5 0) PP(2T (z) (k)) < ( (1) 0) arccos(re.. ) arccos (Ta Arrcos((T(()).)))
2 arccos
1 (1 arccos(r arccos(T 2
arccos(re;,))
where w : /1/2 and Tw,e v 0/ (11 |[@l]) is the cosine of the angle between and u, relative to the inner-product structure induced by Σ. The third line in the above display is a direct application of standard orthant probability formulae for bi-variate Gaussian random variables.
Note that é = e₁ and so Fe... = cos(0)/(ut) = p/y =: r. The result then follows upon invoking
Lemma 2 to estimate ra(k),e and re(k))
13

---

Published as a conference paper at ICLR 2026
Proof of Corollary | Grokking Profile of Vanilla GD). We know that Mk a (1 - (m/mmm/mm a (n1 m2/mm)) m1/m2 = m1 m because the second term dominates. We get
Tk = 1k/Lk
VI+ (B my m1)²
For # O()), have pa for all and can further write
I (uzUs/ua)e
To ensure the test error is Etest (w())) significantly better than chance level, we need rk < r, which translates to
tan 0|: (B8k (11 a (w2m2)/m1)) < 11r1" - = - p1))///0² =etaan
or equivalently, in log space
log T
k k.
with v = /uzll tan 0 m2/m1-
log 1/3"
Note that we have used the fact that log 1/3 logg1 17) x ne, for small E. This gives
k. log
T/8
The case of small [ug| follows a similar argument to analyze the growth rate of arccos(rk) and ultimately get (1/n) log re*
A.2 IMPORTANT LEMMAS
The following lemmas are easily proved via the law of large numbers. Lemma 1. For large n, we have the deterministic approximations
Σ
m2
1 _
(
(m²1)
Dots
(m1/my
(15)
0
where the notation " ignores fluctuations of order Op(n in spectral or L- norm.
The following lemma is direct consequence of the previous via equation Lemma 2. For any deterministic vector R?, it holds that
a ~
a 0 w(k) =
14k w(k) Ew 1x w(k) Ev = pAvi + 80402;
0
3k
1k
(16)
where the notation "~" ignores multiplicatiie factors of order 1/22)
A.3 THEORETICAL ANALYSIS OF ToY MODEL:: THE EQUI-DISTRIBUTION CASE
We now consider version of the toy setting considered in Section with the following modification: the training and testing distributions are equal, consistently with Section 5 We shall establish Theorem which is counterpart of Corollary [ this equi-distributional setting.
Let 2 (O 5) with S diag(1, <), u = en (the first canonical basis vector of IR2), v = (cos 6, sin 0), 0 E [0, тr.. The distribution P (of both training and testing data!) is the law of (x, v) where r has the same distribution as conditioned on the event (z (= v) < O and y = sigu(r u).
The situation is illustrated in Figure
14

---

Published as a conference paper at ICLR 2026
Slow feature x
0.0
-0.5
1.0
-0.5
0.0
0.5
1.0
Fast feature xil
Figure 7: Toy dataset (equi-distributionnal setting). For this illustration, we take 0 = a/4 and 0.1. Training data points are circles, while test data points are stars. The color corresponds the data point's true label y.
vanilla GD
EGD
1.00
Test Accuracy
0.75
0.50
0.25
0.00
102
10a
10*
10
102
10
10a
10º
Iteration (k)
Iteration (k)
Figure 8: Empirical verification of Theorem 2 Solid lines are actual experiments, while lines
correspond to theoretical predictions according to Eqn. (19
Population Moments (Exact 1D Representations) Define the population second moment matrix and cross-moment vector:
Σ = Ep |1:0 1
1 2) c= Eplyo]
W. = --
[l V.
C2
Let ς cot 0/ VE. ó and & the standard normal PDF and CDF, All moments admit absolutely convergent 1D representations:
cos A
Z = ==Pp(<x (x) v) s 50 arccos
cos? 0 sin?
a Z /
2 00
& o(()) ((tt) dt,
d
28
000
(1
nit ((xit)
(6()) (ntt) dt,
Z
((xt)
2
Z 10 tolt) ((xt) dt,
2
C2 == 2 ( o(t) Φ(kt) dt.
b =
tt(t) (kt)) dt,
Z J
Z
With a bit of work, one can show that
x =(e),
A_ 0
M 0(ve)
Vanilla Gradient Descent and Deterministic Equivalent.
Vanilla GD squared loss with fixed
stepsize 77 0 (assuming 7) < 1 1/>maxt (E)): w(k + 1) = w<k) 27 (Ew(k) è).
(0) = (u1. u2)
where Σ, è, and wú are empirical counterparts of E, c, and w. respectively, defined by
a = x X
1 Σ n n
TeeT 6:= TTy
ye w. = WOLs := --1
n
15

---

Published as a conference paper at ICLR 2026
The population recursion (deterministic equivalent, exact in the large-n limit) is
wolk + 1) = no(k) 2n (Eue(k) - c); w(0) = (u1, u2)
Closed form:
wokk = Ww 21-222) ( (w., - -w()).
Let wkk (1i2, vx) T Wa w(0). Define
= 21 1 2na
-2nb
14 = 1-//1 E (0,
2njb
1 -221dd
By Cayley-Hamilton Hamilton Theorem.
r-rk r-r p ak
Bk rk pk
M axI + M,M,
r+ -
T+ - P
we deduce that Ma and are given by
Nk pa= [as(Me U1) + Bx ((1 (m) u1) 27b(1= - -222)
(17)
2% V4 (au(v= - u2) (-2nb(( u1) (1 2nd)( - u2))
(18)
Exact Expression for Test Error. Define the following scalars
cos
with La == v + Ev
cos? 0 + esnn
Pk
LE
Working along the same lines as in the proof of Theorem we obtain that the test error of w(k) is given by
arccos(maxx pk))
Eist (w())) = P(x.))>(y'' (kk < 0)
w(k)
(19)
arccos(p)
From the above analytical formula, see that starts decreasing preciscly when Pk p, i.c.
|vdl |pg| tan 0).
For large 2 ku = (11/n) the fast mode has already converged, so to leading order the drop occurs when
[ux| < T, where T = T(0) == |r.1| tan 0).
Two-Sided Estimate of the Plateau Length. Define the integer k. NU (00) by
k, == minfk 2 0: |Px TT}.
(20)
The slow eigenvector is g- 62 O Ve). There exist constants C(0), 11,1+ such that for all 2 ko.
C- U+ t (v+ (2) I |uxl S Jpu - (1))- = 1 - -2nd.
Hence, when U222 V. UT,
U2 - Vol
U2 Vo
log
< k 2n
log
2n
CT
T/C
Since x_ (E), it follows that for |u2
n.1 > T.
U2
V+l
k,
log
(up to multiplicative constants of order 1 + O(VE)).
T
Putting things together begets the following result.
Theorem 2. If (i) mz(#) S mi mm(9) tan and (ii) U2 > Ccl1, then the test error has a plateau of length O 1. ). If either one of the conditions fails, then there is plateau.
ne
Empirical verification of Theorem 2 is shown in Figure 8
16

---

Published as a conference paper at ICLR 2026
B HYPER-PARAMETERS FOR EXPERIMENTS IN SECTION 5
used in different tasks are listed in Table | In this table, n1 is the number of input bits for the subset parity task, DR is data ratio which shows what fraction of all possible combinations have been used as the training set, LR is the learning rate, WD is weight decay, and BS batch size. In all of the cases ReLU has been used as the activation function.
Table 1: Hyper- for Different Tasks
Task
Setting
π2
DR
Width
LR
WD
BS
k 2
400
50
0.01
10-3 3
Subset Parity
k = 3
100
NA
100
0.042
10-2 2
32
k = 4
100
0.023
10
p = 79
Modular Add./Mult.
p = 97
NA
0.5
512
0.7
10-4
512
p = 127
C RANDOMIZED SVD FOR REDUCING COMPUTATIONAL COST
Performing SVD at each iteration introduces computational overhead. Although in some applications the total wall-clock time to reach high accuracy is not as important as the number iterations (for example, in reinforcement learning it is important to minimize the number of iterations, which trans- lates to interactions), to make EGD more plausible in practice we need to reduce the computational overhead of EGD. One way of doing this is to stop EGD the epoch at which desirable validation accuracy is reached. However, we do not have access to a validation set, we need another method that reduces the per-iteration cost of SVD. Randomized SVD (RSVD) is suitable candidate for doing that. The idea behind RSVD is that we can adopt principles from randomized linear algebra to down project a matrix into lower dimensional space, and then calculate SVD for the smaller matrix (Hakko et al. (01T). We have shown the details of RSVD in Algorithm 3 Our implementation is simplified version of the official scikit-learn scikit-learn developers (002)). As the gradient matrices are not square all SVD decompositions are performed in truncated form which has symmetric cost with respect to matrix dimensions.
Although EGD with exact SVD does not require tuning any hyper-parameter, RSVD has two important parameters: rank and number iterations in the inner loop. We have set the number of iterations what rank is
in the inner loop to 2 in all of the following experiments. In this section, we try the best for each task and what trade-offs we need to consider in tuning the rank. It should be noted that all of the execution times reported in this section have been caleulated on the same Tesla T4 GPU, without any other tasks running and affecting the speed. Details on the experimental setup are provided in Section 5 and the dyper-parameten used is listed in Appendix B
Modular Addition. The first task we are going to study is the modular addition problem. Table 2 shows the number of epochs required to reach 95 percent accuracy, total wall-clock time to reach that accuracy, and the execution time per epoch for each method. As it can be seen, in all of the cases regular SVD outperforms all other methods in the sense of the number of epochs reach the accuracy. However, when we consider the wall-clock time. in two out of three cases RSVD with rank 128 outperforms SVD. These results clearly show that if we care about the wall-clock time, there is an optimal rank that does not delay grokking too much while reducing per-epoch computational cost. Figure | shows how changing the rank of RSVD will change the epoch at which grokking happens. To compare column normalization, which is an approximation of EGD, with EGD and its randomized variants (i.e., EGD both with exact and randomized SVD, with different choices of the approximation rank), see that column normalization achieves 95 percent accuracy later than almost all variants of EGD. However, in the sense of wall-clock time, in two out of three cases column normalization is the best method due minimal computational overhead compared to vanilla SGD.
17

---

Published as a conference paper at ICLR 2026
P
Method
Epochs 95% (1)
Time to ($)
Time/Epoch (ms)
Vanilla SGD
10099
133.3
13.2
EGD (SVD)
222(45.5x)
12.5(10.6x)
56.4
79 EGD (RSVD, 128)
224(45.1x)
11.1(12.0x)
49.7
EGD (RSVD, 1=64)
397(25.4x)
14.0(9.5x)
35.3
Column Norm.
863(11.7x)
15.4(8.6x)
17.8
Vanilla SGD
5392
156.9
29.1
EGD (SVD)
116(46.5x)
10,9(1444x)
94,4
97
EGD IRSVD,r=120)
139(38.8g 8x)
11.5(13.6x)
82.8
EGD (RSVD, r=64)
237(22.7x)
12.2(12.9x)
51.4
Column Norm.
328(16.4x)
(1.3(1)2xx)
30.6
Vanilla SGD
3293
123.8
37.6
EGD (SVD)
62(53.1x)
11.4(10.8x)
184.4
127 EGD (RSVD, 128)
76(43.3x)
10.1(12,2x)
132.5
EGD (RSVD, 1=64)
146(22.5x)
12.1(10.2x)
83.2
Column Norm.
133(24.7X)
5,9(21.0x)
44.1
Table 2: Results on Modular Addition. Comparison of Vanilla SGD, EGD with exact SVD), EGD with randomized SVD (RSVD), and column normalization (a simplified version of EGD) on the modular addition problem, reporting epochs to 95% accuracy, wall-clock time (in seconds). and per- epoch execution time (in milliseconds). The numbers written parentheses are the result of dividing the corresponding value for vanilla SGD by the entry. Therefore, larger numbers in parentheses mean being faster by the written factor compared to vanilla SGD, In all of the entries, lower number outside the parentheses and a larger number in parentheses is better. In the sense of the number of epochs, EGD with exact SVD is the best. In the sense the time to reach 95% accuracy, column normalization is the best, and EGD with RSVD and rank 128 is the second best on average.
Vanilla SGD
FGD (SVD)
EGD (ASVD, rank-128)
EGD (RSVD, rank-64)
Addition Modulo
Addition Modulo 97
Additioon Modulo 127
1.00
0.75
0.50
0.00
10c
10:
10c
10a
10a
10"
10a
10a
10a
10
100
10:
1.0
10a
10*
Epach
Epach
Epoch
Figure 9: Results on Modular Addition for different values of the modulus p. Solid lines correspond test accuracy and broken lines correspond train accuracy. The plots show that using randomized- SVD and reducing its rank prevents EGD from reaching its full acceleration potential. However, when we consider the wall-clock time, randomized SVD with rank 128 seems to be the optimal rank, as it does not delay grokking as much as smaller ranks while increasing per-step speed. Refer to Section 5 for details on the experimental setup and to Appendix B for the pypee-parameters used.
Modular Multiplication. For this problem well, we have compared EGD with exact SVD and its computationally more efficient variant based on randomized SVD (RSVD). Table 3 shows execution time results, and Figure 10 shows learning curves. The conclusion is the same as for modular addition. The simplified version of EGD, which is SVD-free (column normalization), achieves the best wall-clock time. However, if we care about both the number of epochs to reach 95% and the wall-clock time, EGD with randomized-SV with rank 128 is the best overall method, as it reduces the wall-clock time while not significantly increasing the grokking epoch.
Sparse Parity Problem. For the sparse parity problem, we have tried 4 different ranks for RSVD to see the trade-offs more clearly. Interestingly, for all ls values, there is an RSVD setup which achieves 95% accuracy both in fewer epochs and in less wall-clock time compared to EGD with exact SVD. This gives us a clue that optimizing this rank is crucial in practice. The execution time results are listed in Table E and the learning curves shown in Figure 11 For the sparse parity problem,
18

---

Published as a conference paper at ICLR 2026
P
Method
Epochs 95% (f)
Time to 95% (s) (1)
Time/Epoch (ms)
Vanilla SGD
10116
133.5
13.2
EGD (SVD)
221(45.8x)
(1.4(11.7x)
51.7
79
EGD (RSVD, 128)
233(43,4x)
11.6(11.5x)
49.8
EGD (RSVD, 1=64)
(27.3))
12.9(10.3x)
34.9
Column Normalization
814(12.4x)
13.8(9.7x)
17.0
Vanilla SGD
5715
115,4
20.2
EGD (SVD)
107(53,4x)
10.4(11.1x)
97.5
EGD (RSVD, 1=128)
126(45.3x)
9.9(11.6x)
78.6
EGD (RSVD, r=64)
233(24.5x)
12.7(9.1x)
54.7
Column Normalization
328(17.4x)
8.4(13.7x)
25.5
Vanilla SGD
3261
99.8
30.6
EGD (SVD)
65(50.2x)
11.9(8.4x)
183.9
127
EGD (RSVD, 121
1x(44.1x)
9.8(10.2x)
132.0
EGD (RSVD, r=64)
142(23.0x)
11.5(8.7x)
80.8
Column Normalization
142(23.0x)
5.8(17.2x)
41.2
Table 3: Results on Modular Multiplication. Comparison of Vanilla SGD, EGD (with exact SVD), EGD (RSVD), and column normalization simplified version of EGD) on the modular multiplication problem, reporting epochs to 95% accuracy, wall-clock time (in seconds), and per-epoch execution time (in milliseconds). The numbers written parentheses are the result of dividing the corresponding value for vanilla SGD by the entry. Therefore, larger numbers in parentheses mean being faster by the written factor compared to vanilla SGD. In of the entries, lower number outside the parentheses and larger number in parentheses is better. In the sense of the number of epochs, EGD with exact SVD is the best. In the sense the time to reach 95% accuracy, column normalization is the best, and EGD with randomized SVD (RSVD) and rank 128 is the second best average.
Vanilla SGD EGD (SVI)) EGD FRSVD, rank=120) EGO IASVO, rank=64)
Multiplication Modulo 79
Multiplication Module 97
Mulúplication Modulo p 127
1.00
0.75
0.25
0.00
10a
10f
10P
10a
Epoch
Epoch
Epoch
Figure 10: Results Modular Multiplication for different values of the modulus p. for different
values of p. Solid lines correspond to test accuracy and broken lines correspond to train accuracy. The same the modular addition problem, in this problem also randomized SVD (RSVD) with rank 128 seems to be optimal, as it achieves faster per-epoch computation and minimally delays grokking compared EGD with exact SVD. Refer to Section for details the experimental setup and to Appendix B for the saprr-parameters used.
column normalization is not reliable, for A it does not reach 95% accuracy, and in other cases, one of the randomized SVD variants with rank 40 is faster than this method both in terms of the number of epochs and wall-clock time.
Average Over All Tasks. In this section, for each problem consider the best setup for each of the methods: EGD with exact SVD, EGD with randomized SVD (RSVD), and column normalization, and take the average in Table [ These results show that if we optimize the rank for all tasks, on average EGD with RSVD provides the best gain both in the number of epochs achieve 95% accuracy and the wall-clock time.
the gains these provided compared Vanilla SGD. Results are reported
19

---

Published as a conference paper at ICLR 2026
Method
Epochs to 95% (f)
Time 95% (8) -
Time/Epoch (ms)
Vanilla SGD
7873
409.4
52.0
EGD (SVD)
109((2.2x)
28.3(14.5x)
259.5
EGD (RSVD, r=40)
42(187.4x)
9.1(45.0x)
216.7
EGD (RSVD, r=30)
424((1.6x))
70.7(5.8x)
166.7
EGD (RSVD, r=20)
890(8.8x)
137.5(3.0x)
154.5
EGD (RSVD, 1=1
2586(3.0x)
354.5(1.15x,
137.1
Column Normalization
520(15.1x)
40.4(10.1x)
77.8
Vanilla SGD
1784
89.4
50.1
EGD (SVD)
115(15.5x)
36.3(2.5x)
315.6
EGD (RSVD, r=40)
180(9.9x)
39.7(2.2x)
220.7
3
EGD (RSVD, r=30)
76(23.5x)
13.5(6.6x)
178.0
EGD (RSVD, 20)
99(18.0x)
15.2(5.9x)
153.8
EGD (RSVD, 1=1
312(5.7x)
41.4(2.1x)
132.6
Column Normalization
120(14.9x)
8.8(10.1x)
73.8
Vanilla SGD
5273
253.1
48.0
EGD (SVD)
254(20.7x)
50.5(5.0x)
199.0
EGD (RSVD, r=40)
183.0
EGD (RSVD, r=30)
176(30.0x)
27.6(9.2x)
157.0
EGD (RSVD, r=20)
312(16.9x)
45.5(5.6x)
145.9
EGD (RSVD, r=10)
1067(4.9x)
136.5(1.8x)
127.9
Column Normalization
69.9
Table 4: Results Sparse Parity Problem. Comparison of Vanilla SGD, EGD (with exact SVD), EGD with randomized SVD (RSVD), and column normalization (a simplified version EGD) on the modular multiplication problem, reporting epochs to 95time (in milliseconds). The numbers written in parentheses are the result dividing the corresponding value for vanilla SGD by the entry. Therefore, larger numbers in parentheses mean being faster by the written factor compared to vanilla SGD. In all of the entries, lower number outside the parentheses and a larger number in parentheses is better. Interestingly for this problem, in all three cases, there is RSVD setup which leads to faster grokking compared to EGD with exact SVD. The best method both in the sense of number of epochs and wall-clock times is EGD with RSVD with rank 30.
IRSVD,
ECD rank=10)
Parity(nm400, km2) RetU NM(vidth=550
Parityin=100, km3) RetU 100v
ParityIn km4) RetU NNIvidth==100
Accuracy
0.8
0.6
10
10a
10
10a
10" 10:
102
102
Epoch
Epoch
Epoch
Figure 11: Results Sparse Parity Problem. Solid lines correspond to test accuracy and broken lines correspond to train accuracy. Interestingly in this problem, there are some randomized SVD (RSVD) setups which grok in fewer epochs compared to EGD with exact SVD. This task shows that although column normalization is a successful light version of EGD, it is not reliable as it does not lead to increase in test accuracy. Refer to section for details the experimental setup and to Appendix for the hyper-parameter used.
D EGD IN MORE PRACTICAL SCENARIOS AND EMPIRICAL COMPARISON WITH GROKFAST
Grokking has been detected in learning more practical datasets, such as (eCuu et al. 2002), and in modular arithmetic tasks when using more complex architectures, such as 2-layer transformer (Liu et al. 2023 Lee et al. 2024). In this section, analyze whether EGD and EGD
20

---

Published as a conference paper at ICLR 2026
Method
Epochs to 95% f
Time to 95% % (s) f
EGD (SVD)
44.8x
9.8x
EGD (RSVD, r=Best for cach Task)
55.6x
14.7x
Column Normalization
16.9x
13.2x
Table Average of improvement gained by the best setup for each task. For Column Normalization. Sparse Parity Problem with & = 4 is excluded as the test accuracy does not increase at all in this case.
with randomized-SDD can accelerate grokking in these more practical scenarios as well. We also compare our proposed method to Grokfast, which is the strongest baseline available. We have used the same hypper-aramettes and optimizer as reported in the Grokfast paper and used the official codebase provided by the authors to ensure a fair comparison. In all of the following comparisons, we have used the best Grokfast setup, which uses EMA filter gradients instead of memory bank to minimally add memory overhead to the vanilla optimizer. As discussed before, our method does not require any extra memory usage.
In this section, we show that EGD provides meaningful gains in more practical scenarios, when the dataset more complex (MNIST), when the model is more complex than a 2-layer MLP (3-layer MLP and and 2-layer transformer), and when the optimizer is more complex than vanilla SGD (Adam and A).mW).
MNIST Dataset with 3-Layer MLP as the Model and AdamW the Optimizer. It has been shown that with some specific hyper-parametent
a 3-layer ReLU MLP network exhibits grokking when trained on the MNIST dataset (Liu et al.) 2023). We have used the same parameters, optimizer (Adam W (Loshchilov & Hutter 2017)). and model as reported by Omnigrok (Liu et al. (2023) and Grokfast (Lee et al. 2024). We perform EGD on the weights of all three layers of the MLP network. Results are summarized in T EGD and all of its variants based on randomized SVD (RSVD) achieve 85% accuracy in fewer optimization steps than Grokfast. However, EGD with SVD only minimally outperforms Grokfast in the sense of wall-clock time to reach 85% (85 percent is chosen to be consistent with the results reported by Grokfast), but the variants of EGD with RSVD outperform Grokfast by a large margin both in the number of steps and in time to reach 85%. If we look at the learning curves shown in Figures T2 and [3] we see that EGD and its variants grok faster than Grokfast. Also, in Grokfast there is sudden drop and then increase, which we do not see in the EGD plots. Another interesting observation is that the training loss plots are noisy for both Vanilla Adam W and Grokfast. However, EGD has smooth training loss curve.
Method
Steps to 85 % (t) Time to 85 % (8) (1) Time/Epoch (ms) v Final Val. Ace.
Vanilla AdamW
43100
660.0
15.3
89.0%
Grokfast
1800(33.9x)
28.7(23.0x)
16.0
91.0%
EGD(SVD)
700(61.6x)
24.3(27.2x)
34.7
91.9%
EGD(RSVD.r=le 16)
600(71.8x)
11.7(56.4x)
19.5
91.3%
FGD(RSVD,r=8)
500(86.2x)
9.4(70.2x)
18.8
89.6%
EGD(RSVD.=4)
500(86.2x)
9.2(71.7x)
18.4
91.5%
Table 6: Results on MNIST. We have used the same parameters, model, and optimizer suggested by the Grokfast paper. As can be seen, RSVD-based variants of our proposed EGD outperform Grokfast both in the number of optimization steps and in the time to reach 85%. Also, in terms of final validation accuracy, EGD (both with exact and randomized SVD) variants outperform Grokfast.
Modular Arithmetic with a 2-Layer Transformer as the Model and Adam as the Optimizer. Following the setup suggested and used by previous papers (L.iu al.. 2023 Lee al., 2024), we use 2-layer transformer and the Adam optimizer (Adam et | [2014] as the baseline to compare with EGD. We also compare results with Grokfast (Lee et al. 2024). The EGD is only performed on the first MLP layer of transformer. Execution time results are listed in Table 7 The results suggest that EGD with RSVD performs comparably to Grokfast. However, our method does not add any overhead to memory. These results are based comparing EGD with Grokfast, which uses EMA for filtering. However, to calculate EMA, Grokfast requires memorizing a copy of gradients to update. In EGD, no extra memory is needed. The learning curve comparisons are shown in Figures a4d and
21

---

Published as a conference paper at ICLR 2026
Varilla Adami
EGD (preposed method)
Grokfast
Train/Validation Accuracy (MNIST)
Train/Validation Loss (MNIST)
1.0
10a
10a
Accuracy
10"
10-1
2
0.4
10-3
10-4
10
10a
10a
10
10a
10a
10a
10
10
Optimization Step
Optimizatio Stop
Figure 12: Results on MNIST. Solid lines correspond to validation accuracy/loss and broken lines correspond to train accuracy/loss. Here our EGD is implemented via exact SVD. As it can be seen, EGD accelerates grokking more than Grokfast. Also, it does not show any sudden decrease in accuracies or large noise in train loss.
Varilla AdamW
EGD (SVD)
Grokfast
EGD (ASVD, 4ark=4)
EGD (ASVD, rank=8)
EGD (RSVD, rank=16)
Train/Validation Accuracy (MNIST)
Train/Validation Loss (MNIST)
1.0
10a
0.8
10a
0.6
10-2
10-1
10-1
10a
10a
10a
10
102
10a
10a
10
Optinization Step
Optimization Step
Figure 13: Results on MNIST. Solid lines correspond to validation accuracy/loss and broken lines correspond train accuracy/loss, Here, EGD method is implemented via exact SVD and randomized SVD (RSVD) with different ranks. We see that for some ranks, RSVD groks even faster than EGD. All of the RSVD versions have smooth training los curvae in contrast with Grokfast and vanilla AdamW.
15
These curves show that EGD is extremely unstable in this scenario. However, when we use randomized-SVDD, the results become stable and even more stable than Grokfast.
Method
Steps to 95% 4 (f) Time to 95% (s) L (t) Time/Step (ms)
Final Val. Acc.
Vanilla Adam
55780
5198.7
93.2
99.1%
Grokfast
910(61.3x)
114.0(45.6x)
125.3
100%
EGD(SVD)
1170(47.7x)
179.9(28.9) (x).))
153,8
100%
EGD(RSVD.=36)
980(56.9x)
129.3(40.2x)
131.9
100%
Table 7: Results Mudolo Multiplication with p 97 and 2-layer transformer the model. Grokfast outperforms EGD with exact SVD, both in the number of epochs and the time to reach 95% accuracy. However, the difference between the RSVD case and Grokfast is negligible, while EGD with RSVD does not introduce any extra memory usage, which is crucial in large models such transformers.
22

---

Published as a conference paper at ICLR 2026
Adam
EGD (Proposed Method)
Grokfast
TrainWallddation Accuracy (Muttplication Modulo
TrainWalldation Loss (Muttplication Modull 97)
10-
0.5
10-1
10
10
10-30
0.0
10a
102
10a
10
102
10a
10a
10a
10
10
Epoch
Epach
Figure 14: Results on Mudolo Multiplication for p=97. Solid lines correspond to validation accu- racy/loss and broken lines correspond to train accuracy/loss. Here, our EGD method is implemented via SVD, and leads to unstable training for this task and model architecture.
Vanilla Adam
EGD (RSVD, rark=36)
Grokfast
Train/Valddation Accuracy (Multipicaation Modulo 97)
TrainWaldaation Lass (Muttipiication Module 97)
0.8
10-7
0.6
Loss
10-1
10
10a1
0.0
10a
10a
10a
10
102
101
10a
10a
10
10F
Epoch
Epach
Figure 15: Results on Mudolo Multiplication for p = 97. Solid lines correspond to validation accu- racy/loss and broken lines correspond to train accuracy/loss. Here, our EGD method is implemented via randomized SVD (RSVD) with rank = 36. Interestingly, EGD with RSVD here leads more stable training, compared to Grokfast EGD with exact SVD (Figure 14
Summary. In practical scenarios, EGD with randomized SVD (RSVD) leads faster grokking both the number of steps and in wall-clock time without adding any extra memory utilization, which is crucial for large models and optimizers that already use lot of memory to track momentum and velocity of gradients. However, we need to tune the rank of RSVD to achieve the best gain. The results are promising across various tasks, optimizers, and architectures.
E EGD HELPS WIDELY USED OPTIMIZERS HANDLE NON-STATIONARIE
In many real-world scenarios, handling con-stationarity is crucial. We hypothesize that our proposed method EGD, which leads to faster grokking and faster learning of nuances in data, will help optimizers adapt quickly to changes in the distribution of data. In this section, we test how adding EGD on top of four commonly used optimizers can improve their adaptability to changes in the distribution of data. We follow the experiment introduced by Sookar al. (2023) to simulate non- stationarity. The idea is to start training on CIFAR-10 10 (Krizhevsky et al. 2009) and, at fixed intervals, shuffle all of the labels to see how fast the optimizer can recover and increase the training accuracy after shuffling. This experiment has been shown to be effective demonstrating non non-stationarity in reinforcement learning (Castanyer al. 2025
23

---

Shuffling at Epochs Published as a conference paper at ICLR 2026
We have used the same setup and parameters used by Castanyer et al. (2025) for these experiments. The model is a Convolutional Neural Network with 3 convolution layers, max pooling, a fully connected layer, and a classification head. We perform EGD with SVD on all convolution layers and the fully connected layer. To make 3-dimensional convolution layers 2-dimensional, we flatten the last two dimensions. Figure T6 the effect of adding EGD on top of three commonly used optimizers (Adam (Adam et al. 2014), RAdam (Liu et al. 2019). and RMSprop (Tieleman & Hinton] (017)) to enable them to recover faster after shuffling. In all cases, EGD increases the area under the curve, An interesting observation is that, without EGD, RMSprop seems to be more successful compared to the other two optimizers handling con-stationarity. Table (aole/B8itie shows that adding EGD improves handling of non-stationarity for all optimizers. The most significant improvement is observed for the Adam optimizer and when there are points of shuffling.
Number of Shuftling
Method
Adam
RAdam
RMSProp
Vanilla Optimizer
43.0
40.8
60.2
Shuffling at I Epoch
Optimizer + FGD
79.2(1.8x)
70.0(1.7x)
79.5(1.3) .3x)
Vanilla Optimizer
24.7
22.9
31.6
Optimizer + EGD
56.0(2.3x)
47.3(2.1x)
54.7(1.7x)
Vanilla Optimizer
23.0
19.8
40.1
Shuftling at 4 Epochs
Optimizer + EGD
44.0(1.9x)
23.5(1.2x)
45.6(1.1x)
Vanilla Optimizer
14.2
13.3
14.8
Shuftling at Epochs
Optimizer EGD
24.3(1.7x)
22.0(1.6x)
24.0(1.6x)
Table 8: Non-Stationarity Injection Results CIFAR-10. Area under the curve for various optimizers after being subjected to a shift distributions, with and without EGD. The gain from using EGD is most significant for Adam and RAdam, and the largest gain occurs when there are points for shuffling labels.
F DOES EGD CONVERGE TO DIFFERENT SUP-SPACE COMPARED TO VANILLA SGD?
Methods accelerating grokking usually do not converge to the same optimum as vanilla SGD (Lee et al. 2024). To test if that is the case for EGD as well, we have plotted the trajectory of weights while learning arithmetic tasks. The experimental details are the same as what provided in Section B and the dyper-parameterd are listed in Appendix B Figures 17 19 and 21 show the trajectories for Addition, Multiplication, and Sparse Parity problems, respectively, and Figures [8 20 22 show the distance between weights different time steps and the initialization. The following interesting observations can be made from these plots:
Common across all tasks: after the initial big steps, EGD does not move much, but vanilla SGD moves over wider range.
Modulo problems: EGD and vanilla SGD take different and diverging trajectories. Both overshoot in distance from initialization. However, the distance reduced afterwards.
Sparse parity: although initially EGD and SGD trajectories diverge, SGD tries to reach the vicinity of the point at which EGD starts grokking. In terms of distance, both methods show an initial big jump and then maintain an almost fixed distance from the initialization.
Based on these observations, we can say that, except for the sparse parity problem, in the other two tasks EGD and vanilla SGD converge to different regions of the parameter space. In the sparse parity problem, they first diverge, and then vanilla SGD attempts to move toward the region where grokking occurs (i.e., vanilla SGD tries turn toward the epoch at which EGD makes its large jump). The reason behind this behavior, and whether this observation also holds in the randomized case and for other tasks, is left future work.
24

---

Published as a conference paper at ICLR 2026
Vanilia
Optimieer EGD
Adam
RAdam
RMSprop
Train Accurary
42-
Train Accunsy
02
Thain Accuraty
484
Train Train Accunsy
2.8
Figure 16: Non-Stationarity Injection Results on CIFAR-10. In each row, we shuffle the labels the same intervals. From top bottom, the number of shuffling points increases, and therefore recovery becomes more difficult. The results show that RMSprop is more successful in recovering after shift in distribution. Also, in all optimizers and scenarios, adding EGD improved adaptability.
Weiggt (pn79)
Weight trajechory (pm127)
14332
13083
20-
14307
12003
20
1.000
12000
10383
10303
-20
20
PCL
923
PC1
Figure 17: Results on Modular Addition. Trajectory during EGD and Vanilla SGD. EGD and SGD converge to different sub-spaces.
25

---

Published as a conference paper at ICLR 2026
540
Distance from nitllarttion (p-79)
Destance from Initiiaratios (p-977
Dilance from nitfisatiinn (p-127)
120
110
100
1No Wolb
60-
2006 6007 9930
35000 12000 Epech
1006
34000
2000 4232
1000
10000 Opoch
8007
1,000
34000
2002 4000
0000
090e Spoth
19990 12000 10030
Figure 18: Results on Modular Addition. Distance between weights and the initializations. Both methods overshoot, and in the end converge regions with the same and fixed distance from the initialization.
Weight trajectory
trajectory (o-97)
trajectory
16030
10903
114030
12030
12300
12000
-40
4003
1200
4000
4000
-20
2002
-20
2030
-00
Van SCD
Werlll
P03
Figure 19: Results on Modular Multiplication. Trajectory during EGD and Vanilla SGD, EGD and SGD converge to different sub-spaces.
Iniiiizzation
inttiiszation
130
110
126
100
(00 16/s
20-
8000
tpoch
Qooch
Spot
Figure 20: Results Modular Multiplication. Distance between weights and the initializations. Both methods overshoot, and in the end converge to regions with the same and fixed distance from the initialization.
Weight trajectory ()=22
Weig Waiggtttaajestter trajectory (o=3)
Wie g/t trajectory (6m4)
200ce
3750
C22
1500
3250
g
1000
4000
Vini SCC
2000
-2
0
0.0
25
5.2
7.3
12.0
12.5
Figure 21: Results on Sparse Parity Problem. Trajectory during EGD and Vanilla SGD. EGD and SGD first diverge, then EGD turns towards the epoch at which EGD grokked.
26

---

Published as a conference paper at ICLR 2026
5CD
Distance from inw.i (k-
Distence foor Inttilacatinn [K-3]
Didance from aen
10-
1000
2000
4000
2000
6000
10030
250 500
750
1000
1250
15:00
0750
3266
2000
3000
4000
2000
6000
Figure 22: Results on Sparse Parity Problem. Distance between weights and the initializations. EGD overshoots first, but in the end remains in the same distance from initialization. Vanilla SGD does not overshoot and gradually goes to the region with the same distance to initialization as EGD.
G IMPLEMENTATION DETAILS AND PSEUDO-CODES
In this section, provide the pseudo-code based on which obtained the results in this paper. We present detailed algorithms for the three main components our method. The descriptions of these algorithms are as follows:
Training Loop with E This box shows the general training loop with EGD. To reduce the cost, we have access to validation set and threshold for acceptable validation accuracy, we stop using EGD after the epoch at which we attain that validation accuracy. Also, we can choose between RSVD and SVD.
EGD with Optional RSVD 2 This algorithm accepts gradient matrix and performs preconditioning based on SVD or randomized SVD. To get symmetric cost with respect to dimensions of the matrix, we always calculate truncated SVD.
Randomized-SVD 3: A detailed implementation of randomized SVD is shown in this box. We used a simplified version of the algorithm from the scikit-learn implementation (scikit-learn developers, 2024). The main difference between this algorithm and the official version is that we are not using oversampling. Studying the effect of oversampling is left to future work.
Algorithm TRAINING LOoP WITH EGD
Require:
parameters 0, learning rate 7), number epochs Nepochs training data D validation data Dval. threshold X, flag RSVD
1: begin:
use BGd True
2: for epoch = 1... Nepoch do
3: for each batch (X, y) in Drrain do
Compute gradients: G +- VaL X,y) if usD then G + D-EGD(G, RSVD)
4:
EGD preconditioning (Algorithm 2
5:
else
G G- end
No EGD applied
6:
Update parameters: 0 <- OptimizerStep(O,G.. 7: end for
val_ace t- Evaluatee Dva))
8: if val acc > À then
EGD False
Stop EGD updates after threshold is reached
9: end
10: end for
Return: trained parameters 0
27

---

Published as a conference paper at ICLR 2026
Algorithm 2 EGD WITH OPTIONAL RSVD
Input: gradient matrix of size dout dins flag use_ RSVD
2: begin:
3: if use RSVD then
4:
(U,5, VT) RSVD(G)
5: else
6:
(U,S,VT SVD(G)
Truncated SVD
7: end if
8: PSU
Construct pre-conditioning matrix
9: G-PG PG
10: Return: G
Algorithm 3 RRADOMIZED-SVD
1: Input: matrix G of size m x target rank iterations Fiterr seed
2: begin:
3: m,n sGape(G')
4: Q < Normal(n, r, seed)
Sample a Matrix with Random Gaussian entries
5: for i 1. Rater do
6:
Q < QR(QQ)
QR Decomposition
7:
Q QR(G@O) Q)
QR Decomposition
8: end for
9: B --GQ
Project onto low-rank subspace Truncated SVD on small matrix B Return approximate singular vectors and values
( (VD((B
11: Return: Ú
28


Reviews:
[{"review_id": "KT2AX0op7m", "reviewer": "Reviewer_QrdE", "summary": "The paper proposes a buffer-free gradient preconditioning method for accelerating grokking phenomenon, reducing the gap between memorization and generalization from spectral analysis. The authors have provided theoretical results that justify their egalitarian gradient descent (EGD) algorithm assuming two-parameter model. Based on the insights, the algorithm is tested on the typical examples where grokking appears, demonstrating the acceleration.", "questions": "- The box in line 350-355 seems to swallow additional lines after the equation.\n- Some works such as [Kumar et al. (2024)](https://openreview.net/forum?id=vt5mnLVIVo) and [Grokfast](https://arxiv.org/abs/2405.20233) point out that accelerated grokking-induced solution space can be different from the non-accelerated grokking-induced solution space. In other words, the final reaching ground in the parameter space can be differ greatly. Does this proposed method results in the same solution manifold or in the greatly different solution manifold?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The idea is simple and elegantly addressed, with easy theory and clear empirical examples.\n- Experiments show that at least among the given examples, the proposed method effectively reduces the grokking gap.\n- The experimental details are sufficient and easily reproducible.", "weaknesses": "Some points prevent me from giving a higher score. I will raise my score based on the discussion.\n\n- Since the effect of grokking has demonstrated in various other settings than modular arithmetic such as in [Omnigrok](https://arxiv.org/abs/2210.01117) (MNIST, Transformer, Graph CNN, LSTM) or [Kumar et al. (2024)](https://openreview.net/forum?id=vt5mnLVIVo) (MNIST), more experiments can be performed to further justify the clear effectiveness of the proposed EGD.\n- Theoretical justification only exists for toy examples. I believe it should be easy to extend the theoretical results to show that the insights are extendable to larger, more general models such as linear layers or two-layer MLPs.\n- The proposed algorithm seems to be similar to RMSProp. Further discussion between the existing gradient descent variants to the proposed method should be carried on. Moreover, how does the proposed egalitarian gradient descent interacts with “momentum” variable which typically appear in the machine learning context?\n- The connection between FIM, NGD, and the proposed EGD can be further strengthened by showing how do individual FIM-based and NGD-based algorithms behave under grokking environment. This will further justify the proposed method for accelerating grokking, and imply richer connection between FIM and NGD involved in this problem.", "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": "lXAs9TJ4u8", "reviewer": "Reviewer_6NV3", "summary": "This article studies the grokking phenomenon based on new insights on simple classification problems. A theoretical analysis is provided on a binary classification problem, showing how grokking can arise due to ill-conditioning in optimization. Then a modified gradient descent method (egalitarian GD) is proposed to adjust the gradient direction to avoid stagnation during training. Numerical results on sparse parity and modular addition also confirm the effectiveness of the insights.", "questions": "- Can you explain why you choose a different training and test data distribution in Section 3 (fig 5)? This is not a standard machine learning setup.\n- Are you considering the same training and test distributions in the numerical results in Section 5? If so, I would suggest to modify the theory in Section 3 to make it consistent with the setup in Section 5. \n- It is unclear how the G is F in  eq. 11 is estimated from mini-batch samples. This practical aspect should be discussed to better understand the numerical results.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "-\tThe theoretical picture of the grokking phenomenon is interesting. \n-\tThe proposed EGD algorithm works well in practice. \n-\tThe article is well written, with a quite complete literature review.", "weaknesses": "-  The theoretical example is somehow biased and too low-dimensional, which makes it hard to explain practical results.\n- There is no theoretical convergence analysis on the proposed EGD algorithm.", "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": "beTKUJ6pKt", "reviewer": "Reviewer_vvi7", "summary": "This paper introduces Egalitarian Gradient Descent (EGD) as an optimization to reduce and or mitigate the grokking phenomenon. Theoretical rationale is provided as to why grokking occurs within training practices using stochastic gradient descent (SGD). They show that the leading singular directions dominate the optimisation process, which they show plays a role in grokking. Therefore, they propose the EGD, which sets all the singular directions of the gradient update to 1; thus, all singular directions are optimised equally, hence the name Egalitarian. They highlight that this method can be viewed as a simplified version of Grokfast and enjoys the same benefits while not requiring hyperparameter sweeps. The results of the method are shown on the sparse party problem and modular arithmetic tasks.", "questions": "1.  Line `374` states In practice, we turn off EGD and switch it for vanilla (S)GD once we detect grokking has occured $^1$ where the footnote states This is detected by monitoring validation loss.  When does one exactly turn off EGD? At what point do you consider grokking detected? Is it once the model starts to generalise? Does this not introduce a hyperparameter?\n\n2. For the sparsity parity problem, Figure 3. The method does not fully resolve grokking for Parity(n=100, k=3) ReLU NN(width=100). Is there a rationale behind why here the model training and test accuracy do not increase in tandem, as was shown in the modular arithmetic tasks? \n\n3. Can you provide the results of this method on Grokking-Induced MNIST and IMBD dataset as done in [1]\n\n[1] Liu, Z., Michaud, E.J. and Tegmark, M., 2022. Omnigrok: Grokking beyond algorithmic data. arXiv preprint arXiv:2210.01117.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is very well written, reads well, and is self-contained. \n\n- It is positioned very well within the literature exploring grokking and methods to mitigate/reduce grokking. \n\n- It provides a strong theoretical rationale behind the method presented.\n\n- Provides insights into the occurrence of grokking more generally.", "weaknesses": "- In Figures 1 and 2, it is not clear where the train accuracy lines are for the EGD method. Is this due to them going up in tandem, or has this data been omitted? \n\n- Figures 1 and 2 range from -0.5 to 1, resulting in a large portion of the figures being clear; this should be reduced to make the data easier to read. It would also be useful to provide the grid to make reading off the figures easier. \n\n- Increasing the number of tasked benchmarked would increase the impactfulness to the overall community, especially tasks where grokking is induced, such as on MNIST in [1]. \n\n- Although the method is fairly simple, as it requires performing SVD on the gradient to receive $U$ and $V^T$, it would be good to provide the codebase to enable the work to be easily replicated. \n\n- The `Limitations and Future Work` section reads more like future work. The limitations of the work are not clearly represented. This should be more aptly named Future Work. \n\n- The Bibliography is missing links to papers. \n\n- Typo/Spelling Mistake: Line 375 occured should be `occurred`\n\n\n[1] Liu, Z., Michaud, E.J. and Tegmark, M., 2022. Omnigrok: Grokking beyond algorithmic data. arXiv preprint arXiv:2210.01117.", "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": "V4pF2JACvY", "reviewer": "Reviewer_yxV7", "summary": "This paper studies the phenomenon of grokking—where neural networks plateau for long periods in generalization performance before abruptly improving—by connecting the length of the plateau to the spectral properties of the gradient during training. The authors propose Egalitarian Gradient Descent (EGD), a modification to classic gradient descent that normalizes the gradient along all principal directions, theoretically eliminating slow convergence due to ill-conditioned spectra. Extensive empirical evaluations are provided across toy problems and algorithmic tasks, demonstrating that EGD can dramatically reduce (or eliminate) the grokking plateau compared to vanilla SGD and other methods.", "questions": "1. Applicability/Scalability Beyond Toy Tasks: What is the maximal network width, depth, and dataset size on which you have successfully deployed EGD, and how does the runtime/throughput compare with vanilla SGD and other strong baselines? Can you provide concrete experiments or ablations addressing actual wall-clock cost?\n\n 2. Necessity of Full-Rank SVD: Is the full SVD required for each gradient step in all layers, or would it suffice to approximate just the top-$k$ singular values/vectors, perhaps by sketching methods? (Experiments addressing this could open EGD to wider scaling.)\n\n 3. Direct Empirical Comparison to Grokfast: Your only comparative discussion with Grokfast is qualitative. Can you provide synchronized benchmarks (same datasets, architectures, and hyperparameters) for EGD and Grokfast, so the empirical strengths and limitations are evident?\n\n 4. Failure Modes or Adverse Side Effects: Can EGD cause instability, exploding/vanishing gradients, or harm convergence in high-noise regimes or complicated real-world architectures? Any evidence of problems in practical runs would be useful.\n\n 5. Potential for Generalization Analysis in Deep Networks: Is there any theory in progress on whether the spectral-flattening argument for “egalitarian” updates extends robustly to multi-layer, highly nonlinear settings?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Clear theoretical insight and connection to spectrum: The paper gives a transparent analytic treatment for a simplified grokking scenario, precisely quantifying how ill-conditioning in the gradient’s principal directions translates into stagnation in test accuracy (see Section 3; Theorem 1 and Corollary 1). This directly links observable grokking delays to quantifiable optimization issues, moving beyond anecdotal claims.\n\n 2. Simple, principled method: EGD is conceptually straightforward, introduces no free hyperparameters, is optimizer-agnostic, and is presented as a direct, closed-form modification to the gradient for each layer (Section 4). Compared to competing approaches like Grokfast, EGD incorporates no history/memory buffers, making it lightweight for integration.\n\n 3. Sound mathematical formalism: Derivations (e.g., equations (11)-(12), as well as the detailed proofs in Appendix A) are logically consistent, and the singular-value “flattening” interpretation of the update is rigorous.\n\n 4. Strong empirical backing: Results in Figures 1–3 show that EGD consistently accelerates grokking across modular addition, modular multiplication, and sparse parity datasets—tasks where the phenomenon has been robustly documented. EGD typically leads to immediate generalization, while SGD and even “column normalization” lag—this is evident in Figure 1 and Figure 2.\n\n 5. Insightful spectral diagnostics: Figure 4 explicitly illustrates how delays in grokking correspond to highly skewed gradient spectra, making the theoretical contribution tangible and informative for the reader.", "weaknesses": "1. Scalability and Computational Cost: The method requires per-step SVD computation on the gradient matrix for each layer, an operation that grows rapidly with layer width. While the authors claim (Page 7) that randomized or approximate SVDs “work just fine”, there is no benchmark or ablation quantifying the tradeoff, or on what network scales/width SVD ceases to be efficient. This undermines claims about practical usability in modern architectures.\n\n2. Lack of Ablation and Sensitivity Analysis: The experimental section has limited sensitivity checks. For example, EGD is briefly compared to a “column normalization” heuristic, but not to common natural gradient or second-order optimizers (despite Section 4.1’s close connection). There is no ablation quantifying: (i) the effect of running EGD on only top-$k$ directions vs. full-rank, or (ii) across various batch sizes/optimizers. This makes it hard to parse where the real gains come from.\n\n3. Missing Rigorous Quantitative Comparisons with Strong Baselines: Grokfast (Lee et al., 2024) and other recent grokking accelerators are discussed theoretically (Section 4.2), but quantitative experimental comparisons are limited. The paper could benefit from more comprehensive results tables directly benchmarking EGD against these methods across identical tasks.\n\n4. Overstated Generality and Claims Beyond the Evidence: In several places (e.g., Section 6, Conclusion; Abstract; and throughout the Introduction), the paper implies EGD could be universally beneficial and “plug-and-play” for any neural network. The only evidence is for the tightly controlled algorithmic family, and the real computational cost is glossed over. The claims regarding optimizer-agnosticism and hyperparameter-freeness require much more stress-testing.\n\n5. Limited Theoretical Depth on Nonlinear/Deep Networks: Section 4 describes the transition from linear to nonlinear settings, but almost all the theoretical analysis is confined to the analytically tractable linear case, or at best layer-wise argument. No generalization, convergence, or stability claim is rigorously proven for deep, nonlinear architectures—the regime of real interest.", "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)