A Theoretical Framework for Grokking: Interpolation followed by Riemannian Norm Minimisation
Authors: Etienne Boursier
Scott Pesme
Radu-Alexandru Dragomir
Venue: NeurIPS 2025 poster ()
Keywords: grokking; non-convex optimisation; dynamic perturbed systems; theory of neural networks; Riemannian optimisation
Primary area: optimization
TL;DR: We describe the training of overparametrized architectures with small weight decay as a two-phase dynamics. In particular during the second phase, it follows a Riemannian flow of the norm on the interpolation manifold.

Abstract: We study the dynamics of gradient flow with small weight decay on general training losses $F: \mathbb{R}^d \to \mathbb{R}$. Under mild regularity assumptions and assuming convergence of the unregularised gradient flow, we show that the trajectory with weight decay $\lambda$ exhibits a two-phase behaviour as $\lambda \to 0$. During the initial fast phase, the trajectory follows the unregularised gradient flow and converges to a manifold of critical points of $F$. Then, at time of order $1/\lambda$, the trajectory enters a slow drift phase and follows a Riemannian gradient flow minimising the $\ell_2$-norm of the parameters. This purely optimisation-based phenomenon offers a natural explanation for the \textit{grokking} effect observed in deep learning, where the training loss rapidly reaches zero while the test loss plateaus for an extended period before suddenly improving. We argue that this generalisation jump can be attributed to the slow norm reduction induced by weight decay, as explained by our analysis. We validate this mechanism empirically on several synthetic regression tasks.

A Theoretical Framework for Grokking: Interpolation followed by Riemannian Norm Minimisation
Etienne Boursier INRIA
Scott Pesme INRIA
LMO, Université Paris-Saclay
Grenoble, France scott pww.eeerreannia....m
Orsay, France
etienne. boursierriniiaa.fz
Radu-Alexandru Dragomir
Télécom Paris
Institut Polytechnique de Paris
Palaiseau, France
dragomiz@teeecco-pprii..
Abstract
We study the dynamics of gradient flow with small weight decay general training losses F; I2d + R. Under mild regularity assumptions and assuming convergence of the unregularised gradient flow, show that the trajectory with weight decay À exhibits a two-phase behaviour as À - 0. During the initial fast phase, the trajectory follows the unregularised gradient flow and converges to manifold of critical points of F. Then, at time of order 1/A, the trajectory enters a slow drift phase and follows a Riemannian gradient flow minimising the fn-norm of the parameters. This purely optimisation-basel phenomenon offers a natural explanation for the grokking effect observed in deep learning, where the training loss rapidly reaches zero while the test loss plateaus for an extended period before suddenly improving. We argue that this generalisation jump can attributed the slow norm reduction induced by weight decay, as explained by our analysis. We validate this mechanism empirically on several synthetic regression tasks.
Introduction
Strikingly simple algorithms such as gradient methods are a driving force behind the success of deep learning. Nonetheless, their remarkable performance remains mysterious, and a full theoretical understanding is lacking. In particular: (i) convergence to low training loss solutions on non-convex objectives is far from trivial and (ii) unclear why the resulting solutions generalise well [Zhang al., 2017]. These questions training. One such intriguing behaviour is known as the grokking phenomenon, which we explore in this work. Coined by Power et al. [2022], this term describes two-phase pattern in the learning curves: first, the training loss rapidly decreases to zero, while the test loss plateaus at a certain value. This is followed by second phase, where the training loss remains zero, but the test loss steadily decreases, leading to a final improved generalisation performance, as depicted in Figure 1 (left).
accompanied by range of surprising phenomena that arise during
In this paper, we propose novel theoretical perspective to explain this phenomenon. By examining the gradient flow dynamies with weight decay, we show that, in the limit of vanishing weight decay, we can fully describe the trajectory of the model parameters. Specifically, we prove that the training process can be decomposed into two distinct phases. In the first phase, the gradient flow follows the unregularised path, converging to manifold of critical points of the training loss. In the second, the trajectory enters a slow drift phase, where the weights move along this manifold, driven by weight
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

First plase
Second phase
Manifold of critical points of /
W
train and test osses
2.0
Irain Test Iwila
Unregularised gradient flw trajectory
phase:
1.5-
1.5
Second phase: towards lower E/ norm solution
1.0-
1.0
This second phase often leads to grokking!
min MHE MEManifold
0.0
0.0
10
Training time
Figure 1: Gradient flow with small weight decay A. (Left) A typical example of grokking: the training loss rapidly drops to zero, while the test loss plateaus for a long period before eventually decreas- ing- coinciini with a steady drop in the mn-norm of the weights. (Right) Schematic illustration in parameter space _d of the optimisation behaviour described in Theorem 1. The trajectory (*) initially follows the unregularised gradient flow and converges to a manifold of critical points of F (fast dynamics). At time t 1/X, the regularisation term becomes dominant and induces slow drift along this manifold toward a lower f2- fz-norm solution (slow dynamics).
decay, gradually reducing their l norm, as illustrated in Figure (right). We argue that this slow decrease in the weight norms explains the grokking phenomenon, as smaller weight norms are often correlated with better generalisation. To convey the main intuition, we state below an informal version of our result, describing the full trajectory of the gradient flow with small weight decay.
Informal statement of the result. For generic training loss F Rd + IR satisfying some regularity assumptions, we consider the gradient flow w regularised with weight decay: ui = FV (t) (t). Under the assumption that the unregularised gradient flow trajectory is bounded, we prove the following.
Theorem (Main result, informal). As the weight decay parameter is taken to the trajectory w (1) (1) can be seen as composition of two coupled dynamics:
1. (Fast dynamics driven by F given by Proposition 1) In a first phase, the weights follow the unregularised gradient flow and converge to a manifold of critical points of F.
2. (Slow dynamics driven by the weight decay given by Proposition 2) At time t z 1/>, the iterates start slowly drifting along this manifold, following a Riemannian gradient flow that decreases the fo-norm of the weights.
Link with the grokking phenomenon. Note this is purely optimisation result: statistical assumptions are made, and a priori does not imply any improvement in test loss during the slow phase. However, provides natural explanation for the grokking phenomenon. Indeed, in practice, for many deep learning models with random initialisation, gradient flow converges to a global minimiser of the training loss. When this solution generalises poorly-as often the case with large initial weights. in the so-called lazy regime [Chizat et al.. 2019] the subsequent slow drift along the critical manifold, driven by weight decay, decreases the fn-norm of the solution and simplifies it the second phase. Since lower weight norms often correlate with better generalisation [Bach, 2017, Liu et al., 2022c, D' Angelo et al., 2024], this offers a convincing explanation for the delayed improvement in test performance. We discuss various settings where this behaviour is observed in Section 5.
2 Related work
Grokking in experimental works. The term grokking was originally coined by Power et al. [2022], which studied two-layer transformer trained with weight decay on a modular addition task. They observed that the network quickly fits the training data while generalising poorly, fol- lowed much later by sudden transition to near wear-perfect generalisation. Following this work, many studies have investigated modular addition tasks better understand the mechanisms underlying this phenomenon [Nanda et al., 2023, Gromov, 2023]. However, grokking has been observed far beyond this setting. For instance, Barak et al. [2022] showed that training a neural network to learn parities exhibits a similar delayed generalisation pattern. In Liu et al. [2022c], grokking was induced across a broad range of tasks, including image classification and sentiment analysis, by using small datascts, large initialisations, and weight decay. Other settings and architectures where grokking- like

---

behaviour appears include matrix factorisation [Ly et al., 2023] and learning XOR-clustered data with a ReLU network [Xu al., 2023]. Finally, is worth noting that this delayed transition in generalisation was already observed in earlier works, as clearly illustrated in Figure 3 of Chizat and Bach [2020]. More recently, Jeffares and van der Schaar [2025] argued that grokking may not be so central to Deep Learning and may only appear in very specific situations. However, we still believe its a priori counter- intuitivie aspect is worth investigating and might lead to theoretical understandings that go beyond what is currently referred by grokking.
Grokking as the transition between lazy and rich regimes. Several works have framed grokking as the transition between the lazy and rich regimes. The lazy regime, also called the NTK regime, was introduced by Jacot et al. [2018]. It typically arises when the network is trained from large initialisations [Chizat et al., 2019], and corresponds to a setting where zero training loss can be quickly achieved, but often with poor generalisation performance. In contrast, the rich regime (also called the feature learning regime) corresponds to a setting where the network actively learns new internal features during training. In the classification setting, Lyu and Li [2019] show that the rich regime is always attained for homogeneous parameterisations, and similarly, Chizat and Bach [2020] provide analogous result for infinitely wide two-layer networks. In this context, Lyu et al. [2023] and Kumar et al. [2024], followed by Mohamadi et al. [2024], offer a theoretical perspective on grokking as the transition from the lazy regime to the rich regime during training: initially, the predictor quickly converges towards the NTK solution, and later escapes this regime to reach a better generalising solution, driven by the effects of implicit regularisation and/or weight decay.
The role of weight decay. The role of weight decay in grokking remains somewhat debated. While many of the original works exhibiting the phenomenon include weight decay [Power et al., 2022, Liu et al., 2022b], grokking has also been observed without [Chizat and Bach, 2020. Xu et al., 2023], as strongly emphasised by Kumar et al. [2024]. However, as shown in Lyu et al. [2023], the transition tends to occur much later and to be less sharp without weight decay. In this context, weight decay can be interpreted as a factor that triggers or accelerates the transition from the lazy to the rich regime. While grokking can be observed in classification tasks even without weight decay-thanks to the algorithm's 's implicit bias-to to the best of our knowledge, it cannot occur in regression tasks unless weight decay is used. Of particular relevance our work, Liu et al. [2022c] propose an intuitive explanation of grokking that is based on weight decay: during the first phase, the model rapidly converges to poor global minimum; during the second, slower phase, weight decay gradually steers the iterates toward lower-norm solution with better generalisation properties. While appealing, this explanation remains informal and lacks rigorous theoretical support. In this work, we provide a formal analysis of the optimisation dynamics underlying the grokking phenomenon: an initial fast phase leads to convergence toward the solution associated with the lazy regime, followed by a slower second phase that drives convergence toward the solution characteristic of the rich regime.
Drift on the interpolation manifold. Many theoretical works in the machine learning community have studied the training dynamics of gradient methods in overparameterised neural networks, where the set of zero training loss solutions forms a high-dimensional manifold. In this context, leveraging results from dynamical systems theory [Katzenberger, 1990], the work of Li et al. [2021] describes the drift dynamics induced by stochastic noise after stochastic gradient descent (SGD) reaches the manifold. This analysis was further extended by Shalova et al. [2024]. Leveraging a similar stochastic differential framework, Pillaud- Vivien et al. [2022] precisely characterises this drift in the setting of diagonal linear networks, and proves that leads to desirable sparsity guarantees. Much in the spirit of our work, although outside the deep learning context, Fatkullin et al. [2010] derive stochastic differential equations that describe the dynamics of systems with small random perturbations on energy landscapes with manifolds of minima, illustrating how the system first converges to the manifold and then drifts along it.
Our analysis builds upon this framework popularized by Li et al. [2021] and tracing back to Katzen- berger [1990]. Our contribution differs from previous work in both focus and scope. Whereas previous analyses attribute the second, slower learning phase to stochastic effects, we identify a deterministic mechanism -amelll,,, slow drift induced by regularization and link it directly to the grokking phenomenon, connection that, to our knowledge, has not been previously explored. From a technical perspective, our setting is simpler yet enables a more detailed analysis. Rather than appealing to results from Katzenberger [1990] as black box, we provide simpler, self-contained proof building on Falconer [1983]. Furthermore, unlike prior analyses that assume initialization near the interpolation manifold, our framework accommodates arbitrary initialisation. We rigorously
3

---

characterise the initial convergence toward the manifold and the subsequent transition to the drift phase along it, which constitutes the main technical novelty of our analysis.
3 Setting and preliminaries
We consider a loss function F; I + R. which is sufficiently smooth, as stated in Assumption I below. Typical examples include the least square loss over some training dataset, where the parameters to optimise represent the weights of some neural network architecture. For a given À > 0, we define the regularised loss Fx as:
Fx(w) = F(w) + A I 2
Vi E Rd
Initialising the parameters from wo E IRd (independently of A), we then consider the gradient flow w over the regularised loss for any À > 0, as the solution the differential equation
i (t) (t)
and
wh (0) wo-
Gradient flow is the limit dynamics of (stochastic) gradient descent as the learning rate goes to 0. For A = 0, we denote w GF the gradient flow on the unregularised loss:
GF (t) FF (t)) and GF (0) W.-
In the remaining of the paper, we consider the following assumption on the objective function.
Assumption The function F is C on Rd, its third derivative is locally Lipschitz and F is definable in an o-minimal structure. Moreover, the solution w GF to the gradient flow ODE (2) is bounded.
The regularity conditions on F ensure that, for any À 2 O, the gradient flow ODE has a unique solution, which is defined for all t 2 O. This is a consequence of the Picard-Lindelöt theorem and boundedness of the trajectories. Definability in the o-minimal sense guarantees that bounded gradient flow trajectories converge to a limit point [Kurdyka, 1998, Thm. 2]. This is mild assumption satisfied by most functions arising applications, such as polynomials. logarithms,, exponentials, subanalytic functions, and finite combinations of those; see Coste [1999], Bolte et al. [2007] for more details. Overall, Assumption
pretty mild and holds for all architectures (e.g.. neural networks or transformers) that use differentiable activations.
Finally, note that the boundedness assumption on the unregularised flow excludes classification settings where the network can perfectly separate the data, since in such cases the unregularised iterates diverge.
3.1 Stationary manifold and Riemannian flow
Assumption guarantees that the gradient flow converges to a limit point u wFFF = lim GF (f), which is stationary point of F. In typical scenarios of training overparameterised models, stationary points are not isolated, but form continuous sets [Cooper, 2021].
Definition (Definition of the manifold M). We define M to be the largest connected component of VP (0) containing GF where V (0) corresponds the set of stationary points of F.
(0)
Our key assumption is that M forms a smooth manifold, and that it contains only local minimizers (and not saddle points). Additionnally, M are lower bounded by some constant v) > O. Following the terminology of Rebjock and Boumal [2024], this is known as the Morse-Bol Bout property.
impose that the non-zero eigenvalues of the Hessian on
Assumption 2. M is smooth submanifold of Bd of dimension k [d], i.e. for any w E M, = d Also, there exists 1) > such that for any w M, all non-zero eigenvalues of 2 F(w) are lower bounded by v).
As stated in Rebjock and Boumal [2024], for C2 functions this property is equivalent to the Kur- Lojaieeojsse also known as 1998, Bolte et al., 2009]. implies that (i) every point in M is a local minimiser, and (ii) all gradient flow trajectories are locally attracted towards M. This stability property is essential for proving regularity of the flow map & in Section Let us discuss the relevance of Assumption 2 the context of overparameterised machine learning.
(PL) condition locally around M [Kurdyka,

---

Convergence to a local minimiser: our assumption rules out the possibility that w o converges to saddle point of F. This is justified by a large number of works showing that gradient methods avoid saddle points for almost all initialisations. In particular, Lee et al. [2016, 2019] prove it under the assumption that all saddle points of F are strict. Although their result only holds for discrete-time gradient descent, the underlying argument can be extended to gradient flow (the proof relies on the Stable Manifold Theorem for dynamical systems, which also holds in continuous time [Teschl, 2012, $7]).
Morse- Brrrse-Boot/Lojssiwiiie property: this is a common assumption in the analysis of gradient flow dynamics for overparameterised networks [Li et al., 2021, Fatkullin et al., 2010, Shalova et al., 2024]. Note that for general models, the critical set may not form manifold everywhere. However, it is often possible to show that the manifold structure holds locally, i.e., on most of the space, excluding some degenerate points's see Section 5 for examples and Liu et al. [2022a] for a generie result. Moreover, the results derived from such an assumption are generally very representative of empirical observations, as can be seen in Section 5.
Riemannian gradient flow on M. We endow M with the standard Euclidean metric. For any differentiable function h : IRd - I. we denote by grad wilt the Riemannian gradient on M of the function h defined as follows:
grad M RRd w H- PT (Vh(r)),
where PTM(w) (w) the orthogonal projection on the tangent space to M at w. Under Assumption typical properties of smooth manifolds imply that (w) = Ker(V2 F(w)) [see e.g., Boumal, 2023, for a detailed introduction to optimisation manifolds]. Using this notion of Riemannian gradient, we study the Riemannian gradient flow for some objective function h and initialization W.M E M, defined as the curve w satisfying.
(1)) = -radd ((w(()))(((()) and w(0) = WMM-
By construction of the Riemannian gradient, the trajectory of any solution of this ODE necessarily belongs to (0), and therefore to .,, since M is maximal connected component. If h is C² and has compact sublevel sets, Assumptions 1 and 2 guarantee that there exists unique solution to Equation (3) and that it is defined on R..
4 Grokking as two-timescale dynamics
This section states our main results, where we characterise the two-timescale dynamics of the regularised gradient flow (1) in the limit + In Section 4.1 we describe the first phase, the fast dynamics, where w approximates the unregularised gradient flow solution on finite time horizons. Section 4.2 then identifies the second, slow dynamics happening at arbitrarily large time horizons, where - follows the Riemannian flow of the &g norm on the manifold M of stationary points.
4.1 Fast dynamics
first simple observation is that as À + 0, Fx - F uniformly on any compact Rdd From there, seems natural that w should converge, at least pointwise, to GF X + 0. A Grönwall argument indeed allows to characterise the first, fast timescale dynamics given by Proposition below.
Proposition 1. If Assumption holds, thenfor all T 2 a w GF (C"((o, R7, || llso)
Importantly, uniform convergence only holds on finite time intervals of the form [0, T). but is not true on R.. More precisely, grokking is observed when the two limits cannot be exchanged: limx-10+ lim (t) & lim;-+0o lim w> (t) = lim;-+00 w GF (t). In Figure 1, the endpoint of the red arrow corresponds this first limit, while the endpoint of the blue arrow corresponds to the second. This distinction highlights& a key aspect of grokking: the dynamics evolve on two different timescales. Initially, described by Proposition 1, the regularised flow tracks the unregularised one. But at much larger time horizons, the regularised dynamics begin to diverge.
While our results are stated for such global Morse-Boo assumption, they could be directly extended to local Morse-Boo assumptions, provided that the iterates remain in the non-degenerate region.
5

---

4.2 Slow dynamics
The second part of the dynamics is harder to capture, since happens time approaching infinity, when À approaches zero. It is done using theory of singularly perturbed systems. In the following, we associate the unregularised flow function mapping & IR RR + E" satisfying
p(w, f) = w
VFF(( 8)) ds,
V(w,t) E R x R..
Note that 6(w, t) simply corresponds to the solution of the gradient flow of Equation (2) at time t when initialised in When possible to define-is i..., when the gradient flow admits a limit point in Rd--we we define the mapping & as
T (w)
lim p(w, t)-
t-300
Thanks to Assumption I, the unregularised flow initialised in wy admits the limit point 4 (wa), which is necessarily a stationary point of the loss, i... V (wo)) )) = O. Assumption 2 then ensure that the mapping & is defined and C2 on some neighbourhood of M. thanks to result of Falconer [1983].
Lemma 1. If Assumptions and 2 hold, there exists an open neighbourhood U of M such that + is defined and C² on U.
Now the mapping & is well defined on a neighbourhood of M, it can be used to describe the limit of the slow dynamics. For À > O, we let wn : w (t/A). We indeed have to adequately "speed-up" time to capture its behaviour. Notably, w. satisfies the following differential equation:
a (1) 1 VF
w (1)
and
wn-
Our goal in this section is to study the limit function limx- w Intuitively, the VF term in Equation (5) will enforce this limit to stay on the stationary manifold M for any t > 0. The slow dynamics will be shown to approximate the Riemannian flow of the squared Euclidean norm on the stationary manifold M for > 0. This limit flow is defined by T which is the solution of the following differential equation on R4, for the function &
+ |/2]//2
(6) Recall that the Riemannian gradient is grad M (w) = ( (( (w) for any E M. De- noting D the differential of & at w. et al. [2021, Lemma 4.3) proved that for any w E M, "Ker(V=F(w)) D.... i.e., the differential of + at w is given by the orthogonal projection onto the kernel space of the Hessian of F. In consequence, t also satisfies the following differential equation:
()) = -radd (2)))(() and T <
w (I) DD (o)
and now prove our main result, given by Proposition 2.
T (0) 4wwo).
Using this alternative description of w Proposition 2. If Assumptions and 2 hold, then for all T,e > 0, we have w 6 in (I=, R*), II (loc), where uf is the unique solution on R_ of the differential equation (6).
Proposition 2 states that the slow dynamics AA converges uniformly to v as À + 0 on any compact
interval of the form [s, T). Note that excluding 0 from this interval (i.e., E > 0) necessary. Indeed, uniform convergence cannot happen on interval of the form (0, 7). since ū- (~)(0) (0) wi for any À 0 and th (0) (ww)). In particular, Proposition 2 leads to pointwise convergence of w²; we have
lim V (0) wo:
2
lim a (1) = w (L) ift> O.
A=+0
This limit function limx-o i- is non-continuous at O. Indeed the whole fast dynamics, which follows the unregularised flow, happens at that 0 point in the limit À - O. On the other hand, Proposition 2 describes the second phase of the dynamics, starting from the convergence point of the unregularized flow ((wp) - at the rescaled time 0+ and following the Riemannian flow on M.
Note that, once the junction between the slow and fast dynamics is carefully handled via Lemma 2 in Appendix A.2, Proposition can be derived from Fatkullin et al. [2010, Theorem 2.2], which heavily relies on the technical result of Katzenberger [1990]. However, for the sake of completeness and readability, we provide a concise and self-contained proof of Proposition 2, avoiding the use of heavyweight methods from Katzenberger [1990] and relying on weaker assumptions.

---

Sketch of proof. here provide a sketch of proof with the key arguments leading to Proposition 2. Its complete and detailed proof can be found in Appendix A.3. We first define the shifted slow dynamics ÄA for any / 2 0 as (t) = (1 + ((A)), where (()) is the "junction point" between the two dynamics given by Lemma 2 Appendix A.2, and satisfies limx-so (()) = 0. Using Lemma 2, ōA then follows the same differential equation as w with an initial condition now satisfying limx-o (0) (wo) E ... While the dynamics of vA might be hard to control as + 0, it is easier to control the one of ((ö^). Using the chain rule, we indeed have
1 VF((A(()). (0 (D^(t) + À
& (t)
Then using the fact that for any wo in a neighbourhood of M, DT(w) VF(w) = [Li al., 2021, Lemma C.2], this directly rewrites as
c ( (t) DD -0)
Now note that this resembles the differential equation satisfied by w. The two differences being that (i) the initialisation points differ, but lim ë~ (0) u (0); (ii) the time derivative is on ((ō)) rather than 8* directly. To handle the second point, 0* and (()~) obviously converge to the same initialisation point as À + O. One can then use stability of the manifold M, thanks to Assumption 2, to
show that suPre[o.r] | (t) ( (t)]ll converges to 0 as À + 0. This then allows to conclude. O
Characterizing the limit of the Riemannian flow. By monotonicity of its norm, d is obviously bounded over time. Typical optimisation results then guarantee that the limit set of (t) as + co is contained in the set of eritical points of the squared Euclidean norm the manifold M, given by the KKT points of the following constrained problem:
min
WEM
The notion of KKT points indeed extend to smooth manifolds [Bergmann and Herzog, 2019], so that under Assumption 2, the KKT points of Equation (7) are given by the points w² E . satisfying grad (2 (we*) = 0, where we recall grad ful2 is the Riemannian gradient.
We are then able to show that w converges towards the set of KKT points. Note that this does not follow from Proposition 2 alone, as we also need to show that trajectory of wn remains bounded independently of X: we can prove this is true in our case.
Proposition 3. If Assumptions I and 2 hold, then for any sequence (Au) such that Ao
0.
k-+00
the limit points of (lime-soe w> kE)) are included in the KKT points of Equation (7).
While Proposition 3 guarantees that wn gets arbitrarily close to KKT points of Equation (7) as À goes to O, it does not imply that it has the same limit as w It is however guaranteed with the additional assumption that i (l) converges a strict local minimum of the Fuclidean norm on the manifold M.
Proposition 4. Let Assumptions I and 2 hold and, assume additionally that w (t) converges towards a strict local minimum w² of the constrained problem (7). Then lim lim (1) = w*.
When the slow limit dynamics on M converges towards strict local minimum, Proposition 4 guarantees that, for small enough X, w² gets trapped in the vicinity of this local minimum as t - oc, allowing to get perfect characterisation of limx-o lim w (t). In particular, this double limit corresponds to the limit of the slow dynamics wh while the permuted limit (lim, lim w (t)) corresponds to the limit of the fast dynamics GF thanks to Proposition 1.
Role of initialisation scale. The initialisation scale strongly influences the behavior of the unreg- ularised flow and, consequently, the first phase of training under weight decay. This dependence on scale is well-documented in the literature [Chizat et al., 2019]. While a complete theoretical understanding remains open, it is widely accepted that small initialisation scales correspond to the rich regime, in which implicit bias drives the model toward interpolating solutions with smaller weight norms, typically associated with better generalization. In contrast, large initialisation scales give rise to the lazy or NTK regime, where features change little during training. This regime behaves similarly random feature models and tends to produce interpolators with weaker generalisation performance.
In our framework, this distinction has a direct consequence. With a small initialisation scale, the point de (b() reached after the first phase already exhibits small norm ponns-ll corresponding
7

---

to KKT point of Equation (7)-so no second grokking phase occurs, as the system has effectively converged. Conversely, with large initialisation scale, (w=> (0)) retains a large norm, which triggers substantial movement during the second, grokking phase.
The grokking transition is not sudden! In the literature, grokking is often described as "sudden drop" in the validation loss following an extended phase of overfitting. We argue here that this drop only seems sudden when training time is plotted on logarithmic scale, whereas in fact it unfolds over characteristic duration of order 1/2, following a plateau of comparable duration 1 1/X. Indeed, Proposition predicts that the drop takes place within the time interval [e/X, M/ AJ, where E is a small time independent of À (think of E the time it takes for 8 to move very slightly away from wxx ) and M corresponds to the typical time required for x (t) to approach its limit lim; w (t). Prior to this interval, the parameters evolve according to the unregularised flow w GF (t). Let M' denote the typical time it takes for w GF (1) to reach w oo
GF, then the plateau extends over the interval [M', , A.]
As result, on logarithmic time scale, the drop occurs within an interval of length In ln(M/e). while the preceding plateau spans roughly In(=/A) on the same scale. This explains why, as À + O, the drop appears abrupt compared to the plateau in log-scale plots. In contrast, when viewed in linear time, the drop actually extends over a duration comparable to that of the preceding plateau, producing markedly different visual impression.
Comparison with Lyu et al. [2023]. The work most closely related to ours is that of Lyu et al. [2023], which provides a theoretical characterization of the grokking phenomenon the NTK regime-ics i.e., the unregularised flow initialised at large scales- to the rich regime, which typically converges KKT points of Equation (7). However, their analysis does not offer general optimisation-ben bbaed perspective on the phenomenon and, in particular, does not account for the slow drift phase along the solution manifold, which we identify and characterise. Moreover, their setting is more restrictive: it assumes specific network architectures with homogeneous parameterisation and requires large initialisation scales. In contrast, our results hold outside the NTK regime and apply across a broader class of settings. In addition. their theoretical guarantees rely on taking the initial- isation scale to infinity while simultancously letting the regularisation strength tend zero, with both rates polynomially coupled. Their analysis establishes that, for some sufficiently large time f(A), the regularised flow wn approaches KKT points of Equation (7), but it does not provide guarantees about the asymptotic behavior beyond this time. By contrast, Proposition 3 characterises the limit points of the flow w,, offering a stronger and more complete understanding of its long-term dynamics.
transition from
5 Examples and experiments
Linear regression. Let F(w)= = ||x - v1I2 with X E Rnndd and assume that min F = 0. The problem is convex and the set of critical points is the affine subspace M = {w Xw y}: Assumption is satisfied globally.
It is well known that unregularised gradient flow Gr converges to (Mo(w)), the projection of the initial point on A [Lemaire, 1996, Gunasekar et al., 2018]. Then, since M is convex, the Riemannian flow on M necessarily converges to the minimal Eg norm solution w" = X y, where X denotes the pseudo-inverse. Those two points are different (unless wo = 0), which leads to grokking, as wn is expected to have better generalization properties than (w [Bartlett et al., 2020]. In this setting, the trajectories of w can be computed explicitly to illustrate the two-timescale dynamics; see Appendix C.
Matrix completion. This is prototypical non-convex problem which is amenable to theoretical analysis. The goal is to recover a matrix M E gaxm, which is assumed to be low-rank, from subset of observed entries in Ω C [m]] x [m], by solving
F(U, V) = Euser ((VV - -MM))
min UEBeX" V
(8)
where r is the target rank. If the rank of the ground truth M is known, one can set r accordingly. However, the true rank is often unknown. An alternative approach is to use overparameterisation and choose r much higher than needed. Although in this case F has many minimizers, our results indicate that the gradient flow trajectories (2) for small tend converge towards low-rank solutions.
8

---

More precisely, we analyse the extreme overparameterised setting when = m + n. In Appendix C, we show that the set M¹ of stationary points of F which are nonsingular matrices forms a manifold. Provided that unregularised gradient flow converges to a nonsingular point, we can apply our results locally. These results state that, the second, slow phase of the dynamics, the trajectories minimise | P + [VIl on M Recall that for a given matrix M E BRxx,, we have
|M = min (IUTT) + IIVI)),
UVT=M 2
where ||M is the nuclear norm [Srebro al., 2004, Lemma 1]. Since minimising the nuclear norm promotes low-rank solutions, this indicates a drift toward low-rank matrices during the slow phase of the dynamics. In Appendix C, we study the more general class of matrix sensing problems and discuss an important technical subtlety: the set MM forms a manifold only after excluding degenerate points. Handling those singularities highly non-trivial and remains an open direction for future work.
Figure 2 below empirically confirms this grokking for matrix completion. We here randomly generate a rank 3 ground truth matrix M* E R20xx0 with non-zero singular values of, o 0z. We randomly sample 50% of the entries to define the observed entries Ω. We then perform gradient descent with weight decay parameter À = 10-3 and stepsize 7 10-2 on the loss F(U, V) defined in Equation (8) and where the weights U, V E _20x10 are initialised with i.id. standard Gaussian entries. We then track the training loss, the unmasked test loss || UV [[,, the weight norms |||| = || + [[V and singular values of the reconstruction matrix UV Additional experimental details can be found in Appendix D.I.
15
Train Test
Unngelaated
400
Singular valtees of
10-
300
Loss
[W9]
02
03
100
5-
10-1
10a
10a
10-1
Training time L
Training time
Low-rank matrix completion. (Left): Grokking phenomenon: the training loss drops quickly to zero, while the test loss remains high for an extended period before eventually improv- ing- coineiiigg with a decrease in the norm of the weights ||wl|2 = || + [V (Right): Singular values of UVT over time. Each line corresponds to the i-th singular value of UV The singular values rapidly converge to large positive values at time t s 1. However, as grokking starts around time = 102, all but three begin decay towards zero. The remaining three approach the true singular values of, o22 and oz.
Figure 2:
Explaining the observed grokking phenomenon. At time / = O. the weights are randomly initialised and the training loss is high. Initially, the regularised and unregularised weights follow the same trajec- tory. and the training loss quickly drops to zero: the regularised iterates converge to the same solution as the unregularised gradient flow. This carly solution has a high norm, large singular values, and poor generalisation performance. As training continues, around time to decrease. By Rs 10a, the parameters have drifted to new solution that still achieves zero training loss but has a much lower norm and actually coincides with the low rank ground truth matrix M*
1/,, the weight norms begin
Two-layer ReLU network. Although our theoretical framework does not allow for non-smooth architectures such as ReI.U networks, we illustrate Figure that similar grokking dynamics can be observed in this case. We train a two-layer ReL.U network of the form (w) = -i= M2 ReI.U(r, b,), with weights w (u, v. b) where the outer layer is RT the inner weight E R and bias b E "". The teacher function f is a sum of 3 ReLUs and is represented in dotted light blue in Figure 3.
We generate training dataset of n 10 points by sampling x; uniformly in [-2, 2] and computing y: = fxx).. These training points are shown as black crosses in Figure 3. We train the student network
2n Σ (fw(xz) - )1)² using gradient descent with weight decay X = 10-3 for T = 106 iterations and small step size. The initial weights are independently sampled from a Gaussian of variance 4. At each iteration, we record the train loss, the norm of the weights, as well as the test loss over fixed test dataset (plotted Figure 3, left).
with m 100 by minimising the squared loss F(w) = 1

---

0
12
10
Train
1200
1.5
1000
10-
20 1.00
800
X
600
Teacher function Learnt functiio fue
30 10'
167
-2
-1
Training time
Figure 3: Two-layer ReLU network trained with gradient descent and small weight decay, (Left): Grokking phenomenon: the training loss drops quickly to zero, while the test loss remains high for an extended period before eventually improving-coinen norm. (Right): Snapshots of the network's 's prediction function at various training times. The ground truth teacher function (a sum of three ReLUs) shown in dotted light blue, and the training samples are shown as black crosses.
with slow, steady decrease in the weight
Explaining the observed grokking phenomenon. At time & = O, the weights are randomly initialised and the training loss is high. By f/2 1, the training loss has dropped to nearly zero, and the iterates closely approximate the solution that would be obtained by unregularised gradient flow, this solution does have a low norm and generalises poorly. Subsequently, around time & 1/2, the weight norms begin to decrease, and by f3 ~ 10 they have drifted to a zero training loss solution which has much lower fg-norm and which generalises much better. Such solutions are believed to have small number of "kinks" [Savarese et al., 2019, Parhi and Nowak, 2021, Boursier and Flammarion, 2023], as observed in Figure (far right plot).
Diagonal Linear Networks. We also study-both as an application of Theorem I and numer- ically- the architecture of diagonal neural networks in Appendices C and D.2, which serve a toy problem for neural network training dynamics. In that case grokking promotes sparse estimators.
Conclusion
This work presents rigorous and general optimisation based description of the grokking phenomenon two-timescale process. In the fast initial phase, parameters evolve according the unregularised flow until reaching a stationary manifold. In the slower second phase, they follow the Riemannian gradient flow of the norm constrained to this manifold. Grokking naturally emerges from gradual simplicity bias: starting from a poorly generalising solution recovered by unregularised gradient flow, the slow phase driven by weight decay gradually simplifies the model by reducing its norm, ultimately leading to better generalisation.
While prior work has extensively analysed the first phase via the implicit bias of optimisation algorithms, the second phase-nort norm minimisation constrained to the interpolation manifold received little attention. Our framework highlights the critical role of this phase and motivates further study of optimisation dynamics on interpolation manifolds.
Large initialisations (NTK regime) are known to yield poor generalisation [Chizat et al., 2019, Liu al., 2022c], while small initialisations (rich regime) can lead to slow convergence or convergence to suboptimal solutions for the training loss [Boursier and Flammarion, 2024a,b). Grokking may offer a desirable compromise, achieving fast convergence to an interpolating solution while retaining strong generalisation.
Note that our analysis is derived in the asymptotic regime À -+ since this allows for tractable analysis. Extending the theory to a fixed À is considerably harder, that said, in Appendix E we offer heuristic analysis regarding how small À needs be for grokking to emerge. Also note that our analysis can easily be extended to other types of regularisations. In particular, we believe empirical observations reported for training with Sharpness-Awar
Minimization [Andriushchenko and Flammarion, 2022, p. 7] may also be interpreted through the lens of grokking, albeit driven by SAM-style regularisation rather than standard weight decay. Lastly, our analysis focuses on regression settings with bounded dynamics. In classification tasks, by contrast, the stationary manifold lies "at infinity" once interpolation achieved. Extending our approach such settings remains an open and promising direction for future work, likely requiring techniques tailored to classification losses.
10

---

Acknowledgments and Disclosure of Funding
E. Boursier would like to extend special thanks to Ranko Lazic for his insightful discussions, which were instrumental in initiating this project. S. Pesme would like to thank P. Quinton for carefully reading the paper and providing valuable feedback. R. Dragomir is chair holder from the Hi! Paris interdisciplinary research center composed of Institut Polytechnique de Paris (IP Paris) and HEC Paris.
References
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-awar mini- mization. In International conference on machine learning, pages 639-668. PMLR, 2022.
Francis Bach. Breaking the curse of dimensionality with convex neural networks, Journal of Machine Learning Research, 18(19) 2017.
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 Neural Information Processing Systems, 35:21750- 217644, 2022.
Peter I. Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 11 117(48)330033-3000 2020.
Ronny Bergmann and Roland Herzog. Intrinsic formulation of kkt conditions and constraint qualifi- cations on smooth manifolds. SIAM Journal on Optimization, 29(4):2423-2444, 2019.
Jérôme Bolte, Aris Danilidis, and Adrian Lewis. Tame functions are semismooth. Mathematical Programming, 2007.
Jérôme Bolte, Aris Danillidis, Olivier Ley, and Laurent Mazet. Characterizations of Lojasiewicz inequalities: Subgradient flows, talweg, convexity. Transactions of the American Mathematical Society, (62((6))33111-3336) December 2009.
Nicolas Boumal. An introduction to optimization on smooth manifolds. Cambridge University Press, 2023.
Etienne Boursier and Nicolas Flammarion. Penalising the biases in norm regularisation enforces sparsity. Advances in Neural Information Processing Systems, 36:57795-57824, 2023.
Etienne Boursier and Nicolas Flammarion. Early alignment in two-layer networks training is two-edged sword. arXiv preprint 10X1 2024a.
Etienne Boursier and Nicolas Flammarion. Simplicity bias and optimization threshold in two-layer relu networks. arXiv preprint arXiv:2410.023488 2024b.
Emmanuel J Candes. The restricted isometry property and its implications for compressed sensing. Comptes rendus. Mathematique, 346(9-10):58-599. 2008
Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on learning theory, pages 1305-1338. PMLR, 2020.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Yaim Cooper. Global minima of overparameterrieed neural networks. SIAM Journa on Mathematics of Data Science, 3(2):676-611, January 2021.
Michel Coste. An introduction to o-minimal geometry. RAAG Notes, Institut de Recherche Mathé- matiques de Rennes, 1999.
Francesco D' A'nego,, Maksym Andriushchenko, Aditya Vardhan Varre, and Nicolas Flammarion. Why do we need weight decay modern deep learning? Advances Processing Systems, 37:23191-23221 2024.
Neural Information

---

KJ Falconer. Differentiation of the limit mapping in dynamical system. Journal of the London Mathematical Society, 2(2):356372. 1983.
Ibrahim Fatkullin, Gregor Kovacic, and Eric Vanden Mijnden. Reduced dynamics of stochastically perturbed gradient flows. Communications in Mathematical Sciences, 8(2):4994461, 2010.
Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.02679 2023.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, pages 1832-1841. PMLR, 2018.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and
generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Alan Jeffares and Mihaela van der Schaar. Position: Not all explanations for deep learning phenomena equally valuable. In Forty-second International Conference on Machine Learning Position Paper Track, 2025.
Gary Shon Katzenberger. Solutions of a stochastic differential equation forced onto a manifold by large drift. The University of Wisconsin-Madison, 1990.
Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the tran- sition from lazy to rich training dynamics. In The Twelfth International Conference on Learning Representations, 2024. URL
net/ foot//ffmmmm forun?id=vt5mmll viVoo..
Krzysztof Kurdyka. On gradients of functions definable in o-minimal structures. Ann. Inst. Fourier (Grenoble), 48(3):769-788) 763, 1998.
Jason D. Lee, Max Simchowitz, Michael I. Jordan, and Benjamin Recht. Gradient descent only con- verges to minimizers. In 29th Annual Conference on Learning Theory, volume 49 of Proceedings of Machine Learning Research. PMLR, 2016.
Jason D. Lee, Ioannis Panageas, Georgios Piliouras, Max Simchowitz, Michael I. Jordan, and Benjamin Recht, First-order methods almost always avoid strict saddle points. Mathematical Programming. 176(1-2):31 1-337, February 2019.
B Lemaire. An asymptotical variational principle associated with the steepest descent method for convex function. Journal of Convex Analysis, 3:63-70, 1996.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss?-a mathematical framework. arXiv preprint arXiv:21100.m 10..21111 2021.
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over- wer-parameterize non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 59:85- 116, 2022a.
Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. Advances in Neural Information Processing Systems, 35:34651-34663, 2022b.
Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data.
arXiv preprint arXiv:2210.0111 2022c.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arv: Xiv:11966005800 2019.
Kaifeng Lyu, Jikai Jin, Zhiyuan Li, Simon S Du, Jason Lee, and Wei Hu. Dichotomy of early and late phase implicit biases can provably induce grokking. arXiv preprint arXiv:2311.18811 2023.
Mohamad Amin Mohamadi, Zhiyuan Li, Lei Wu, and Danica J Sutherland. Why do you grok? theoretical analysis of grokking modular addition. arXiv preprint arXiv:2407.1233 12332, 2024.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint 1:X1v22301.05217,, 2023.
12

---

Rahul Parhi and Robert D Nowak. Banach space representer theorems for neural networks and ridge splines. Journal of Machine Learning Research, 22(43):1-40, 2021.
Scott William Pesme, Deep learning theory through the lens of diagonal linear networks. Technical report, EPFL, 2024.
Loucas Pillaud- Vivien, Julien Reygner, and Nicolas Flammarion. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. In Conference on Learning Theory, pages 2127-2159. PMLR, 2022.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- cralization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.022 1.2177, 2022.
Quentin Rebjock and Nicolas Boumal. Fast convergence to non-isolated minima: four equivalent conditions for с functions. Mathematical Programming, pages 1-49, 2024.
Pedro Savarese, Itay Evron, Daniel Soudry, and Nathan Srebro. How do infinite width bounded norm networks look in function space? In Conference on Learning Theory, pages 2667-2690. PMLR, 2019.
Anna Shalova, André Schlichting, and Mark Peletier. Singular-limit analysis of gradient descent with noise injection. arXiv preprint arXiv:2404. 122933, 2024.
Nathan Srebro, Jason Rennie, and Tommi Jaakkola. moximum-marghh matrix factorization. In L. Saul, Weiss, and L. Bottou, editors, Advances in Neural Information Processing Systems, volume 17. MIT Press, 2004.
Gerald Teschl. Ordinary differential equations and dynamical systems, volume 140. American Mathematical Soc., 2012.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametriize models. In Conference on Learning Theory, pages 3635-3673. PMLR, 2020.
Zhiwei Xu, Yutong Wang, Spencer Frei, Gal Vardi, and Wei Hu. Benign overfitting and grokking in relu networks for xor cluster data. arXiv preprint urXiv:2310.025.. 2023.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In ICLR, 2017.
13

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: informal version of our main results provided in the introduction. Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings. is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitations of each of our assumptions are discussed after stating these
assumptions. Additional limitations are also discussed in the experiments and discussion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-tex system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: all assumptions are clearly stated and theoretical results are proved in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: experimental details are given in Section 5. Additional details are given in Appendix D for reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
15

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: code is provided in the supplementary material and will be publicly released (via the author github) after publication.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
http:/gdddsss////eeeee ceiemnissssoooioooiiiioooooooooi oneioonooooi olicyyyy)))lo for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. cc/ wwc//ullc///aeeeccceeeceecceelee http//uudee//eeeee Ceaees/CooSSubmieec cessinnoolccyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: see 4.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: the impact of randomness very small in our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
16

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [No]
Justification: the experiments are small and run on personal laptop.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: : //nuurirs.... cc http:/pllc/tttisssuiit
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: no socictal impact
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
17

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: the paper poses no such risks
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: the paper does not use existing assets
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
18

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: the paper does not release new assets
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
19

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: the core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA
that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https : h/tp://ies...//eeeeeeeeee httpe/rris.cc//ooeeeeeeeeeeeeeee / for what should or should be described.
20

---

Appendix
Table of Contents
A Main proofs
21
A.1 Proof of Proposition
21
A.2 Junction between fast and slow dynamics
22
A.3 Proof of Proposition 2
23
A.4 Proof Proposition
24
A.5 Proof of Proposition 4
25
B Auxiliary proofs
27
Proof Lemma
.r..............................
27
B.2
Minimality of F on neighbourhool
28
B.3
Alternative equation for "7"
29
B.4
Bounding the trajectories
29
B.5
Proof Lemma
31
B.6
Striet Minimality
32
Applications
33
D Additional experiments
36
D.1 Additional Experimental Details..
36
D.2 Diagonal linear networks.
36
E Heuristic analysis on how small À needs to be for grokking to emerge
37
Main proofs
A.1 Proof of Proposition 1
Proposition If Asssupption holds, then all T 2 wn w GF (C) R1, |l- (lss)
X-0
Proof. We first need to restrict the dynamics of (L) to some compact of Rd Without loss of generality we can inflate the compact Assumption 1, so that we can assume there some compact K of RRd such that for all
0, B(w GF (L), 1) C K, where ((w GF (1),1) 1) is the ball of radius centered at w GF (t). We also define for any À > 0, Tx inf E R4 w (t) K).
Thanks to the continuity of 2 F. VF is Lipschittz on K, i.e.,
[VF(() VF('')l s clww w'll
for any w, w e K.
We then derive the following inequalities for all f E Tx).
[| w GF))
i - v GF (s)dss
OF (s)) vF) Xw (s)ds
VF (s) VF (s)) ds l (s)lds
$ e f [[~\(s) - w GP (F(()) ds + At sup ||]]
wEK
21

---

The integral form of Grönwall inequality also called Grönwall- Bellmman inequality then leads to the following inequality for any [0, Tx): 1 (t)
w GF (D) S
Aee
sup
||
(9)
wek
In particular for a fixed 7 E 4. there is a X > O small enough such that for any X S X²,
Vt [0, 2], Ate sup
|| 1.
wek
Given the definition of Tx, Equation (9) and the fact that USER_ B(wGF(t) C K, this then implies that for any À < X" 7x > T. In particular, for any À S Xr, Equation (9) becomes
(w) GF (t) - XTe sup ||||_ wEK
sup
rE(o.T]
Proposition then directly follows.
A.2 Junction between fast and slow dynamics
In the limit À - 0, the whole fast dynamics described by Proposition is crushed into the point of the slow timescale. While Sections 4.1 and 4.2 respectively provide descriptions of the fast and slow timescale dynamics, one needs to control the junction of these two dynamics. This junction made possible by Lemma 2 below, as well as timeshift argument in the proof of Proposition 2. Lemma 2. If Assumption holds, there exists a function t(A) such that lim t(A) = and lim i (t(A) = lim 20 GF (t). X
t-
Lemma 2 states that for well chosen timepoint (()), which is of order - A In()., the slow dynamics solution w will be close to the limit of the unregularised flow at that timepoint. This result will then be key in showing that 1o A admits right limit in which is given by lim w G.
Note that this order - (n) for t(A) is necessary: a smaller value of t(A) would not allow enough time for the flow to approach the limit of the unregularised gradient flow; while larger value would correspond to time where the regularised flow significantly drifted from the unregularised one.
This time - A (A) can indeed be seen, in the fast timescale, the point where the dynamics transition from mimicking the unregularised flow, to the slow dynamics that minimises the regularisation term within some manifold.
Proof. Equation (9) the proof of Proposition yields that for any [x, (1)()) - GF (t) S Ae suPwek [[o]]. We can then define ((A) that for any S min(Tx, (AA),
-A
> and observe
2c
[1) -w GF (t)ll $ =
sup |
wek
VV)nn(A)
sup
||/]
2c
wek
Since VAln(A) - O, we have for X small enough that the above term is smaller than 1. In particular
((A) 1
for \ small enough, Tx 2 X From there, Equation (9) yields at e(A) for any small enough À > 0
+(A)
GF t(2), lim GF (t)ll
l (t(A))
lim w (t)l (t(A)) w GF t-000
I||
VAln(A)
||/|| + mmll GF t(x)
lim w (F(())]-
sup
2c
weK
t-000
Now note that our choice of E(A) is such that both the first and second term in the last inequality converge to 0 - indeed, t(A)e 2(A) proof.
-VAln(2)
O and 1(2)
A too which concludes the
2c
"The more classical form of Grönwall inequality cannot be directly used here, since ||w>(1) w GF (t]l might be oon-differentiabee at some points.
22

---

A.3 Proof of Proposition 2
The main element prove Proposition 2 is Lemma 4 below. We first need to state another auxiliary lemma, given by Lemma 3 below, and proven in Appendix B.5.
Lemma 3. Consider Assumptions and 2. Let )x20 be family of solutions of the following ODE for any À >> 0 and 2
* (0=== 1 (F(
If also u² (0) + ug E M, then there exists a neighbourhood of M such that & is C² on U and for every 0, there exists a family of neighbourhoods (Us)sso of M such that
1. Ug U C for any & <,;;
2. there exists A(e) > 0 such that for any S M()), the trajectory (1))20 is contained in Uzi
3. Neso U= M.
Note that the existence of a neighbourhood U in Lemma 3 is guaranteed by Lemma 1. We can now state our key lemma.
Lemma 4. Consider Assumptions and 2. Let (u )x>a be a family of solutions of the following ODE any X > O and 2 0,
(1) VF(
also (0) uo M, then w² converges uniformly, as À 0, on any interval of the form [0, T] to the function u defined as the solution of the following ODE:
v1 = uo
ú(t) -D (u(t))
Proof. First consider neighbourhood U of M such that t is C2 on U and Lemma holds. From there thanks Lemma we can assume that À chosen small enough so that (t) U for any R.. We can compute the time derivative of obu" (((((t)) (t)), using the chain rule for any
t E R, and À small enough:
(1²())+ -
(x² = -D
Then using the fact that for any E U, (D) (F(w) [Li et al., 2021, Lemma C.2).
(())(t)) DD ())))) (t).
(10)
It now remains to show that u (t) M as À 0. guarantee that o and u² (t) have the same limit, which will be done using Lemma
Thanks to Lemma 3, we can consider a family of neighbourhoods (Ue) of M and function A: R + Ri satisfying Lemma 3. As we can always take a smaller choice for any value A(E), we can also choose the function À so that
it is non-decreasing:
lim A(E) = O.
For the remaining of proof, define the function w + DD$ w and take e(X) = inf X < A(E). We consider in the following À small enough so that s(A) is defined and finite. Since lime-so X(e) = 0, e(A) > 0 for any O. The function & non-increasing, so admits a limit Moreover for any 8 0, e((()) < 8 by definition, that lim E()) 0.
23

---

Thanks to Lemma 3. (u (t)) is contained in UG())- By monotonicity of ||ul2,, the trajectory (u(+))c20 is bounded. Moreover, the trajectory of u (t) is also bounded independently of À thanks to Lemma 7. We can thus consider compact K of IRd such that for any 2 0 and X 0, a(t) K and u(t) E K.
Recall that 4 is the identity function on M and is C² on U. In consequence, we have? suPweknU.ss) [b((w) wll2 x O.
Summing over Equation (10) yields for any t 0:
u (t) - L. H (( + (t) (t) + (0)
Since & is C² on U, H is c-Lipschitz on K, for some e > 0. A comparison with u then yields for any t t20
[|1(()) - u(t)l - / (H)( ((((s))) (1) 2 (t)) + (0) - uo
c / (1)())) - uss) u())(ls
sup WE (0
(w()) - w|l + x n (0)) - 240 ll
Similarly to the proof Proposition 1, an integral form of Grönwall inequality yields for any t 2
(t) - u(t)ll <
sup WEKNUE
[w(w)) - w| + (()(n(()) - moll ex.
(11)
Noting that the multiplicative term SUPWeKKUEa [o(w) w|l (f uoll goes 0 as
goes to O allows to conclude on the uniform convergence of u² to on T0.
Proposition 2. If Assumptions and hold, then for all T, = > 0, we have w
T in
X-10
((, (loo), where w is the unique solution on R_ of the differential equation (6).
Proof. Consider the shifted slow dynamics 8² for any t 2 0 as (t) = w + t(A)) with t(A) given by Lemma 2. Using Lemma 2, T then follows the following ODE:
- VF(((())
(t) = -0 -
with an initial condition satisfying limx-vo (0) =(wo).
We can then direct apply Lemma 4 above on أ² which yields that ग. converges uniformly on any interval of the form [0, T') u
Proposition 2 then obtained by observing that lim t(A) 0, so that for any e > and X small enough such that t(A) < e, it holds for any t E [s. T)
| (t) T (t) = | (t- t(A)) w (t)l
$ (t t(x)) w +())) + (t) (t (((t-((A)) f(A))lI
The first term converges 0 uniformly for t [E by uniform convergence of e towards o and the second term also goes uniformly to 0 by (uniform) continuity of d on the considered interval.
A.4 Proof of Proposition 3
Proposition 3. If |Assumptions and 2 hold, then for any sequence (Nk)en such that A -> 0,
k0000
the limit points of (lim w> (t)) included in the KKT points of Equation (7).
This is direct consequence of the fact that & is the identity M. locally Lipschitz and limx-oo e()) 0.
24

---

Proof. By definition [see e.g., Bergmann and Herzog, 2019], the KKT points of Equation (7) are the points wt E .M satisfying
grad 0. Since grad (()r()) KKT points of Equation (7) are the points w e M satisfying
wt Ker(V2F(w"))
(12)
Thanks to Lemma 7, the trajectories (w> are all bounded and w == lim w² exists for any 0. In particular, this limit stationary point of the regularised loss Fx, i.c.
V (F + Aw = 0.
In particular, w = VV
Let (As)keN be sequence in R4 such that λ (w )x Thanks to Lemma 3, w² E M. Moreover, the equality w = Au (w first implies that ( II O (Ax). Rebjock and Boumal [2024, Proposition 2.8] then also implies that
0. Let w" be limit point of the sequence
k-
d(u M) ()+).
Moreover, noting Wi arg min wEMBlu -wl, a Taylor expansion yields
1 VVF((u))(W - -wk) + of 1 FF(wx)
Ilu wk
(w =
+
A A
d
dk
d( d(w
22 F(F(w)) wx) + ol
A
d
F(Wg))( - wK) + o(1).
A6
The equality w
V F 2+) then implies for the subsequence k, associated the limit point
d
w" that
Akn
lim F(W()))))
Wkn w
(13)
Akn
Let uk PKor(V22(wx))
wid we
Note that
(²F(wx) Ak wk F2F(wu)uR d [[F2F(wx)us 2 n/|uxl1,
and
thanks to Assumption In consequence, (ux)s bounded. In particular, admits an adherence point tho Rdd
Since Wke, 200 wn and 2 continuous, F[ - - 0. So that Equation (13)
n-000
becomes
F2F(w" Uxx =
In particular, it yields that w E Im I )). By symmetry of the Hessian, Im(VVF Ker (cr(V2F(w"))) = so that w² E (er((2F((w)) + i.e., it satisfies the KKT conditions of Equation (7).
A.5 Proof of Proposition 4
Proposition 4. Let Assumptions 1 and 2 hold and, assume additionally that w²()) converges towards a strict local minimum w* of the constrained problem (7). Then lim lim, w² = w².
Proof. For this proof, denote w² = lim, d (1) and F F(w*). w² is strict local minimum of the Euclidean norm on M. Moreover using the Morse Bott property [Assumption and see
25

---

Rebjock and Boumal, 2024], we can consider an arbitrarily small > 0 such that the following
conditions simultaneously hold in (0,, 5) for some 0:
Vu E MM n B(w", (20), w # w² = ||0011 < ||011
Ve E B(w", 6), F(w) - F* == " "dww, d(w,M)²
(14)
4
Yu E B(w", .)) ((F(()(()))))))) - F*) 2 [VF 2 (F(w) - F*).
First observe that the strict minimality assumption implies, through Lemma 8, that there exists Eo > 0
(independent of A) such that for a small enough
inf Fx(w)= > F* + A |l211 0 .5)
+ AEu-
(15)
2
Now fix an arbitrarily small 8' E 00, Let to R such that 11 w01l S 4.4 By pointwise
convergence of (to) to / (ta), we then have that for À > 0 small enough, (to) w'Il || < 2.
Without loss of generality, we can even choose tg large enough and À small enough so that for some arbitrarily fixed ε > 0,
(w) S F" +E. From there, we define for this proof T3 = inftl 2 to (1) & 8')). Similarly to the proof of Lemma 7, we have for any E taa 23
(16)
dF (t))
<nn( (1)) F*) X(Ri クリヒ n²)
dt
where R₁ B( .8") || /F(W)]] and Rt = supweB(ww .²) |/0l|]. Again, a Grönwall argu- ment implies that for any 1 It1.
F1 (1)) < F* + se 11- ( + >
R2
R
2
In particular, if we define = min(" to In(A) 2, we have similarly to the proof of Lemma 7 that for any & [t):
S -2² -
10 w n²(1)) - w w²(
for some constant C' independent of and X. In particular, we can choose and small enough so that this quantity is smaller than 8' It then implies that 1' < Ta 2
and
X
R1
R
F (w) SF
A(
+
+
+
2
From there, by monotonicity of the loss, for any 2 t':
R
R²
Fx(w (t)) < F A
+ + 2
1
Also note that Ri +
n²
->
| ||
So we can choose and e small enough so that
2
80
2
R
R
||
2 +e<
+ Eu-
2
From there, the previous inequality implies that for any t 2 t',
(u (t)) F + X 1 112
+ EO).
2
By continuity, Equation (15) then implies that for any 2 t', w w²)()) E B(w*, .8)
To summarise, we have shown that for any small enough 8 > 0, there exists (7) such that for any À (0), lim w> (t) B(w", 5).
This means that limx- lim w (t) w*, which proves Proposition 4.
26

---

B Auxiliary proofs
B.1 Proof of Lemma 1
Lemma 1. If Assumppiosss 1 and 2 hold, there exists an open neighbourhood U of M such that 4 is defined and C² on U.
Proof. First, restrict ourselves bounded open set B of Rd and consider d(:;)) as function B + R for any t.
The main point of the proof is show that MM B is geometrically stable in the sense of Falconer [1983], i.e., that there exists a neighbourhood (in B) U of afMn B, 0 and h < 1 such that for any E U,
d(o(w, Mn B) & bd(w,,Mn
and
o(w, t) E U,
where d(w, B) infREMMB w| - z/12.*
Let MNB B. The Morse-Bott property (Assumption 2) implies that there exists neighbourhood (xx) C B of x, such that F satisfies the Polyak- jojyakktojssieiiiz (PL) inequality with constant 3. thanks the equivalences between both conditions [Rebjock and Boumal, 2024]
IVF(w) 2 (F(ww) F(x)) Vw E Ux)).
(17)
In the following, we define F* = F(x), which is the value of F on the manifold M (the definition does not depend the choice of 2:).
particular, there is some do (z:) > 0 such that (ix, do(x)) c U(x). Thanks to Rebjock and Boumal [2024, Propositions 2.3 and 2.8, Remark 2.10], we can even choose do (z) small enough so that there are some a; B3 such that
(F SB/F(u F+ for any w E ((x, M) F(w) F addlw,)) any w E B(x,60).
(18)
(19)
By boundedness of B, a and 3 can be chosen independently of a E EMnD here.
Now let E Bxx, (x)) and define Tww) = inf 2 0 ((w, 1) & B(x, (o(x))) Necessarily,
T(w)= > and for any E [0, T (w=),, Equation (17) applies to olu, t), so that for any E [0, Tw)))
d(F(d(w, ()) F*) -I|VF((( 1))1
dt
< (n((((( t)) F*).
So that, for any t (w,T(w)):
Fó((w, t)) F² < (F(w)) - n"
Moreover for any E [o, ((w)), Equation (18) also applies, so that
l t) - w/ S / |VF( 8) s2)ls
<
3V(((w) - -F+)e- ns
23
< ((F(w)-F)
By continuity of F, let 6(x) > 0 small enough so that for any w E B(x, 6(x)), 5(x) 23 ((F(w)-FF < 2
da(ar)
The previous inequality then implies that for any E B(x, 8(x))
'The definition of Falconer [1983] is stated differently but is implied by our notion of geometric stability. when taking f(w)= 6(w,t). t).
27

---

and E [0, T(w)):
11( x|| s 1|w - x|| + [1 1) wl
- <6(x) + 28(x) v(F(w) F*)
"
60(x)
2
In partieular, for any w E ((zr,6())) T(w) oo and (ww, B(x, )o(x)) for any / 2 0.
Also, note that [ t) x I (o(m) and (æx dv)))) implies that Mn d(o(w, M). From there, Equation (19) implies for any 2 0 and w (xx, 6(x)):
(F(((u w. t)) F+)
" (F(ww) - -F*)
80 M B)?
In particular, for any & 2 0, we can choose a sufficiently large I such that d(p(w, (),M) B) s kd(w, M B).
M such that Ulerk Blxi, ((x))) 2 Mn B. We then define U(B) = Use(K] B(x, 0(xr)), which is also finite covering of M n B. We then take large enough such that for any E U(B), d(o(w, 1), B) & kd(w,M M n B) for k < min:e[(]] 6(x.)
By compactness of M B, there is a finite family (II)εE|K)
max:e[k] 8(z,) In particular, our choice of l is such that, for f w ,(w,t), (B) invariant by f. Indeed, note that for any w E UB),
(((w), B) s Mal(w,Mn B)
minsejk] 0(x;) max 6(z;)
d( d(w, M
< min 6(zz).
kE(K
In other words, there is Mn B such that (w] x:l < min K] 8(2r)) Moreover since Uie[k] B(x,, (xx))) is covering of MnB, there is [K] such that | 2z 2 minje[K] d(x;) and by triangle inequality:
|| - x,, < 8(1;),
ie.. w U(B).
Since (B) is invariant by f and < 1, M.MB is geometrically stable, so that we can apply Falconer [1983, Theorem 6.3 and Theorem 5.1]. It then implies that & is C² U(B). Taking increasing sequence of open bounded sets B, covering whole Réd we can then define U Un U(B,) and conclude that T is C² on U.
B.2 Minimality of F on neighbourhood
Lemma 5. Consider Assumptions I and 2. Let U be an open neighbourhood of M such that & is continuous on U, then necessarily for any U M, F(w) suP.EM F(x).
Proof. By definition of M, is constant on M so that supsem F(x) infrEM F(x) F*. Moreover ((M) -M and by continuity, 4(U) C M.
28

---

For any w E U M, note that
((w) - - 1.0 FF(((w, 1))dt,
-- F(w) - = F(d(w, 0)) F(((())) = [[VF(((w di
The first equality is the definition of ((w), while the second comes from deriving over time the function + F()(w,t)) and noting that ((w) E M.
In particular, for any U \ M. (ww) # 0, so that the second integral is also non-zero. In particular, F(w)> > F for any w E U M.
B.3 Alternative equation for i
Lemma 6. If d'Assumption holds, the unique solution of Equation (6) also corresponds to the unique solution of the following equation:
(t) = DD (1) (w (t)) and T (0) = (ww))-
Proof. This is direct consequence of the two following equalities for any M:
grad M (2( ((())
D (w.)
The first one is consequence of the definition of the Riemannian gradient and the fact that (w( Ker(V* F(w)) [see e.g., Boumal, 2023, Theorem 3.15 with FF being the local defining function of M]. The second one given by et al. [2021, Lemma 4.3].
B.4 Bounding the trajectories
Lemma 7. If Asssumpttins 1 and 2 hold, there exists compact K of Rd such that for any X > 0 and t 2 0, w (t) E K. In particular, lim4-+oo w²(()) exists for any À >
Proof. Similarly to the proof of Lemma 1, we can consider a neighbourhood U of M where the PL inequality holds:
F || 2 nF F*) Ve E U. Additionally, we can consider > 0 such that ((((w)), c U and
[VF(w))l < 3/ F(w)=
for any w E B($(wo), 6).
1) For some fixed E > 0, Lemma 2 then implies there is a X > 0 and times t(A) such that for any À (0) both hold
{(1) (wo))l < and (x
t(A) ( F(((wo)) Se.
||w
Moreover, the proof of Lemma also implies that there is some compact K₃ of RRd such that for any < A and f ((A)
x (^ E Ki.
2) Now fix À < A and define TA = inf(t 2 A(A) - () & ((o(wo), 5)}. By continuity, T3 The PL. inequality then applies for any E (()) Tx):
((A)
X
FV (t)) 17 (F((((())) F*).
In particular, this allows derive the following inequalities for any E ((A)) Tx):
d.F, d.F (F))
IVV (w (t))l
dt
< -
< (()(((w>(t)) - XR
< (nn(((((00))) - F*) + X(Ri + T R²),
2
29

---

where R1 we ((oo)))) || | and R = supwen [ In particular, Grön-
wall inequality implies that for any t E (T))
((A)
(t) - F* v ( 1(2)
R1 + R²)
)) FM
2(
A
nf 2
S Ee nit 1i1--11-00 4p0
+ X(
R1
+ R²).
(20)
Define 1' min( ((A) + -In(A)
Tx). Using Equation (18) for any E ( ((()) (a)
:1:
w w²(1) w²^( {(1)
(F)(w
(1) w² ((2) )
(F) + f
[s)
1022
(() - Frds Aln(A)
S
(1)0(w))1 + +
From there, Equation (20) yields for any t E (1(A) .t7:
x
(1)( w w ((A))
Voe 2' 3(t' ((2),
Ri
1
Aln(A)
S
A
+ R2) R2a
(1(((wo)) + 5)
- BBVE (VAIn(A) Aln(A) (1)((w)) [b) | 6);
for some constant C, which is independent of both and X. In particular, we can choose and X small enough, that []w w ( I|| & for any (*AA) ('). By definition, this implies that
< Ty, i.e., for any t ((()) f1), wt (t) B({(wo).
3) Since t' <T,, 6' (()) X
In(N) by definition and
R1 +
1
E ((w
A
R2 +s).
2
By monotonicity of the objective, then have for any / 11:
R
Fx (t)) < F* X
R² +
(21)
Now define 7 inf {12 H(A) (1) & E by continuity that for any E [t', TAl:
Since M minimizes F on U, Equation (21) implies
R +
m (t)ll
R + 2
In particular, for K₂ = [00 2R: + 122 + 3s) and any & [e' TxJ. (² E K2. By continuity and compactness, inf F(w) > F thanks to Lemma 5. In consequence, we can choose small enough so that Equation (21) implies that for any 1 [1', T.).
F(w wE(OU)
inf F(w).
Assume that 7x Since w (T) K2, the previous inequality implies by continuity that (Tx) & OU, i... w (Tx) E Ú. This however contradicts the definition of Tx, so that Tx co. In particular for any 1 2 l'. w (1 K2.
To summarize, we have showed that there exists small enough X1, such that for any À S X
30

---

1. w (t) is included in some compact K of IR fort - 40:
2. w (1) is included in (()(w))) 6) E (*A) , I');
3. w (L) included some compact K2 IR for 2 1';
where K and K2 are both independent of A. In particular, there exists a compact K of IR independent of A such that for any À S X*, the trajectory of (w (t))r0 is included in K.
For 2 X², we directly have by monotonicity of the objective that for any
W (011
1 | (0)
1 x (F() F(w
[0)0 + + (())
2
so that the trajectory (w (t))e20 is also included in a compact independent of
As a consequence, the definability assumption of Fx along with the boundedness implies that lom w (t) exists thanks to Kurdyka [1998, Theorem 2].
B.5 Proof of Lemma 3
Lemma 3. Consider Assumptions and 2. Let (u=) be family of solutions of the following ODE any X and
(tt== - 1FF(()) (F))(tt).
If also u² (0) + un E M, then there exists a neighbourhood of M such that & is C2 on U and for every O, there exists family of neighbourhoods (Uz)e>o of M such that
1. Us Ugr U for any <';;
2. there exists A(2) 0 such that for any X(E), the trajectory (u (t)))2) is contained in Usi
3. Neso M.
Proof. We consider the neighbourhood U defined as in Lemma 7. By definition, F is constant on the manifold M and denote its value F*, i.e., F supwEM F(w) = inf weM F(w).
For any E > O, we define Ug as Ug = {w E U F(w) < F + }} and show that it satisfies these three conditions. By continuity of F, Ug is neighbourhood of M and the first condition is obviously satisfied.
Thanks to Lemma 5, for any w U \M, F(w) > F*. This implies the third condition,
n n U-=MM.
<>0
The arguments of Lemma 7 extend to any family of solutions (u ) satisfying the assumptions of Lemma 3. In consequence,
consider a compact K of Rdd such that for any 2 0 and À > u (1) E K. From there, note again that inf inf F(w) > F*
Since n² w0) - un E M, we can then choose X(e) > small enough, so that for any À E (0, X(e)),
A F(u>(0) + 1 (0) < min(F" + inf F(w)). WOUN))
By monotonicity of the objective over time, (() + Atll> min(F*
e, inf (e(we( F(w)) for any 2 0. Since (1) is continuous, implies that (1) U. for any 2 0, which concludes the proof of Lemma 3.
O
31

---

B.6 Striet Minimality
Lemma 8. Under the same assumptions than Proposition with w - w ((),, there exists
6* > 0 such that for any E (0,8*), there exists E > 0 and X > 0 such that for any X E (0,
inf Fx(w) > F + 1www 1 BB( 6)
+ Ae
(22)
2
Proof. Let be such that for any MM B(w", 26*), + w = || > | and such that Equation (14) holds. Now let E (0,8*) We now fix A arbitrarily small, choose A (0, X*) and define
inf
||012 1100112-
4
WEM
f 22
By strict minimality, compactness and continuity, 0.
Let w @B(w", 5). We can decompose
e, where E arg min w'EM || w Necessarily, ||a]|| = d(wr, M)) I 8 and || w'll = In particular, we also have 26 2 | w' 2 5 - ||a|| From there, using the quadratic growth property (Equation 14):
(x(w) 2 F(w) + 2 (1|ull - |mll)²
2 F* + ! d(w, d(w M)² + >
2/2 (11ll12 - 2 ([][ell
4
2
2 Fx(w)) + L dww M)² M||="||+200 + 2
[1lul] 1|0" 112).
4
Let c(e, 8) = min : | (
There are two cases.
1) Either ||a|| S 0), in which case 25 2 ||u w'
> $12 that by definition of
X (1|u|12 w|w= 11²)
Fx(w) 2 Fx(w") A(I1w*1 28)
2 Fx(w)) X(1 28)c(E, + 4s
2 (ww")) + λe.
2) Or ||||| 2 c(e, 8), in which case we simply have, also using that ||a||| - 8:
(x(w) 2 (x(=)) + " ||/ A(Iw=11 + |
2 Bx(") + 7 5) - A(|1=*1 + 26)5.
In particular, choosing X small enough depending on and 8 - we have for any À S X that
" ccc, 6)2 >(Ilw>' + 25)0 2 Ae
So that in both cases, (x( > F + A || + Ae
32

---

C
Applications
In this section, we provide additional details to the examples discussed in Section 5, and specify how our theoretical results can be applied in various settings.
Linear regression. We consider F(w)= ||Xw قll with X Rn and I d; assume for simplicity that X is full rank. In this setting. the dynamics can be computed explicitely to illustrate result.
Denote the solution of minimal &g norm with w' x y. where X is the Moore-Penross, pseudoin- verse of X. The problem is convex and the critical set of F is the affine subspace M = =* + Ker which is a manifold: Assumption 2 satisfied.
Consider the singular value decomposition X = UEV T where U & Raxd V E Rddd are orthogonal and Σ diag(a ......---.. .a) with an+1
σa O. We make the change of coordinates V and notice that in this basis the minimum norm solution 2n' V w= is of the form
2" 2,,...22,, O, O,...,0). 0). Then, we can compute the trajectory of the gradient flow on Fx, initialized at (0) V wi::
for Sis n.
=²(11= zi +e - (z:(0) 2/,000
7,00
with A,00
o
(23)
o X
for (n + 1) S d,
21) e=->(((()))
(24)
Eq. (23) describes the dynamics along the directions orthogonal to M, and Eq. (24) along those parallel to M. When X + O, the first is much faster than the second. In the first phase, the iterates converge to (z)-00, ..... 2,.00 2n+1 (0)..... (0) ((()))))))))))
(2i.... 2n 2n+1 (0), ................................ =0))); this is the limit of unregularised gradient flow 2 GF (which is also here the projection of the initial point onto M). the second phase, the iterates converge slowly towards the mimimum norm solution (zi, 2 0, .......) 0).
Diagonal linear networks (DLNs). DLNs serve toy example to understand the influence of the architecture on the training dynamics of neural networks [Pesme, 2024]. The corresponding optimization problem writes
min
||X (wi w2)
=1l2.
(w).
where O denotes the componentwise product, and X E Rx is the feature matrix with n S. d, which we assume to full rank. It is usually convenient to perform a rotation of the coordinates and rewrite the problem as
F(u, v) = [ - vll.. (2.))RR2d
min
where w², x2 denotes the componentwise square. The critical set of F' is composed of the couples (u, v) satisfying
X (X - --)) - =0 0, X (X(u² -²)) 0²)- 0
(25)
This set has singularities for points who have null coordinates; we exclude those problematic points, we can show that it is manifold.
Proposition 5. The set M = VF-1(0) n (I is a smooth manifold of dimension 2d n.
Proof. Let (ū. D) M Denote W a neighborhood of (ū, 8) such that U (R+)2d, The function H 2d - Rd with (u, X (X e22) -y) -y) is local defining function for M in the sense that for (u, v) E W. we have (u, E M - H v) = O.
The differential of H (ū, 8) is the linear map satisfying for (Au, Av) Rd
DH (), 2x T Au Av).
33

---

It is clear that, since all coordinates of (ū, B) are nonzero, the map (Au, Av) + + Au 8 Av a surjection on Re and therefore rank ((, (ū, 7)) = rank (XT X) n. This proves that MM is manifold of dimension 2d [Boumal, 2023, $3.2].
Because of the singular points in M, the function does not satisfy Assumption 2 globally. However, results can still be applied locally: the paragraph below for details.
Noting that, for a vector w E IR", we have
||ll = min u.ERR
||u|2 + 1|l12 subject to u² x² w,
conclude that in the second, slow phase of the dynamics, the Riemannian gradient flow which minimizes ||u] + ||u|| 2 MM tends to drift towards solutions of low & norm.
Low-rank matrix sensing/completion. with m S n² and y IRm. For given target rank S n, the matrix sensing problem is
Let A: 8" + 1 be linear map on symmetric matrices
min F(W) = [(A(WW - - WEReXP
(26)
A typical example is symmetric matrix completion, where the goal is to recover an unknown matrix
M E Rxx from subset of observed entries with coefficients l2 E {1...n] J:: the objective function writes F(W)- Einnen ((WWT), )ig Mj)² Note that the asymmetric case presented
in Section 5, Equation (8), can also be written as symmetric matrix completion problem, by setting
U V
W
E (n+m))x
and choosing a new mask that selects only the off-diagonal blocks of WW
Usually, one looks for a low-rank solution to Problem (26), by setting r to small value. Here, we choose to rather study the overparameterised setting where n. Our results imply that, though we do not explicitly impose a low rank structure, the gradient flow trajectories W are driven towards a low-rank solution in the second phase of the dynamics.
Similarly to the example of diagonal linear networks, we show that, the overparameterised setting, the critical set of F is manifold if we exclude singular matrices.
Proposition 6. Let F be the matrix sensing function defined in (26), and denote 12 maxx the set of
invertible matrices of size X If the set M VFF1 (0) (0)] NRRK gaxa is smooth manifold.
Proof. The gradient F is
VF(W) 4A" (A(WWT)- - W. bW R"*"
where A R >> Sa is the adjoint of A.
Let W M and let neighborhood of W such that LI C RRxx For W U. W is invertible and we have W E M if and only if A" (AA(WW))))) > O. The function H(W) A (A(WWT)-- is therefore local defining function for M Its differential at W satisfies for U e Raxn
DH [U] A A +UWT).
Since W is invertible, the map @ U + WUT + UW is a surjection from Raxa onto S": indeed, note that for any S", we have (U) Z with = 1 Therefore, the rank of (H(W) is equal to the rank A A for any W e M which proves that M is a smooth manifold.
Dealing with singularities. In the last two examples, the set VF- has singular points, and so Assumption 2 does not hold globally. However, we showed that it holds on "most of the space", as there exists negligible set S such that M = VF (0) S is smooth manifold.
Our results can still be applied locally, assuming that the unregularised gradient flow u GF converges to a point GF E M Indeed, in that case there exists a neighborhood of GF such that v VF-1(0) nu is included in M Then, the Morse-Bot property holds in this neighborhood.
34

---

Consider then the Riemannian gradient flow w of the & norm on M initialized at GF For any time horizon T' such that the trajectory of wD stays l on the interval [0. T7, we can restrict our analysis to this local region, where our assumptions are satisfied. We can then invoke Proposition 2 to conclude that u converges w uniformly on intervals of the form [c,
However, key limitation arises when analyzing the long-time behavior: the results characterizing the limit points Proposition 4) do apply wD converges to singular point outside M This situation can occur, as singular points might correspond to points that minimize the E2 norm on M (e.g., sparse vectors for diagonal networks, or low-rank matrices for matrix sensing). Establishing convergence of w to such singular points remains an open and challenging problem, which we leave for future work.
In summary, our results capture the grokking dynamics near nonsingular points M", but do not yet account for potential convergence toward singular points, which represents important open challenge.
35

---

D Additional experiments
D.1 Additional Experimental Details
In all our figures and to align with the continuous-time analysis, training iterations refers to the rescaled "training time" & = yk, where k is the number of gradient steps and r the gradient descent stepsize. We run gradient descent 10" iterations for Figure 2 and 10 iterations for Figure
D.2 Diagonal linear networks.
f2 10a
13 10a
Trin and best losses
Train
6.3-
M
gall
Juhle
-0.5 0.0 0.5
1.0
-1.0
-6.5
0.0
0.5
3.0
1.0
Training time
Teacher
function 6d
Figure 4: Gradient flow with small weight decay À on a two-layer diagonal linear network. Regression dataset. (Left): Empirical observation of the grokking behaviour. The training loss rapidly drops to zero, while the test loss remains flat for an extended period before eventually decreasing. This transition coincides with a slow but steady decrease in the f2-norm of the weights. (Three plots on the right): Visualisation of the model predictions throughout training. The dotted light blue curve represents the teacher function, and the crosses indicate the training data. Snapshots of the model's prediction function at various training times (shown in increasing colour intensity) illustrate how generalisation is affected before and after the transition at t 1/A.
Experimental setup (Figure 4). We train a two-layer diagonal linear network of the form f fx) (u O v, (x))) where w (u, v) Id2d and @ denotes element-wise multiplication, on ID toy dataset. The input R is mapped to high-dimensional feature space via the feature map p(x) 1, cos (==) .....coos
- sin
xz 2
sin
pd with dj 30. The teacher 2
mm and is shown dotted light blue curve in Figure 4. The training dataset consists of n = 12 input-output pairs (æ,, yi), where Hy are sampled uniformly in [-1,1] and y/ = fzz)). These training points are shown as crosses in Figure 4. We optimise the squared loss F(w) = 2n Σx (y: - (x;)) using gradient descent with weight decay À = 10 Finally, the initial weights are sampled from a centered Gaussian of variance 0.1.
function is sparse Fourier series (x) = + cos GFz
sin
Explaining the observed grokking phenomenon. At time 1 O, the weights are randomly initialised and the training loss is high. By 122 = 102, the training loss has dropped to nearly zero, and the iterates closely approximate the solution that would be obtained by unregularised gradient flow. This solution is fully characterised by the implicit regularisation result of [Woodworth et al., 2020], and it does not have a low norm. 5 Subsequently, around time i = 1/A, the weight norms begin to decrease, and by l3 2 105, they converge to the minimum-norm solution (u*.v*) = arg (w.v)) ||u]l + ||0l1 straightforward calculation shows that the elementwise product 3* == u* C v" solves the problem arg (fx,), ||3 l1- This is L m-miiiiisationn problem, which (under RIP conditions) is known to recover the sparsest solution [Candes, 2008], explaining the zero test loss after the grokking phenomenon.
'One could also reach the solution observed at time l3 = 10 without using weight decay by employing a much smaller initialisation scale (Woodworth et al., 2020J. but the cost of longer training time.
36

---

E Heuristic analysis on how small À needs to be for grokking to emerge
Note that our theoretical results are derived in the asymptotic regime À - 0, since this setting allows for tractable and general analysis. Extending the theory to obtain explicit results for fixed À > 0 is considerably more challenging. That said, we can offer some intuition regarding how small needs be for grokking to emerge.
Grokking depends on a clear separation between two phases: an initial phase where the iterates converge and stagnate at the solution of the unregularised gradient flow, and second phase driven by weight decay, during which test performance improves. For grokking to be observable, the regularised gradient flow should approach the unregularised limit before weight decay begins to significantly influence the dynamics.
To formalize this intuition, we can define two characteristic times: tGF, the convergence time of the unregularized gradient flow, measured as the time which the gradient norm substantially decreases relative to its initial value; and second time IWD of order 1/X, associated with the onset of the regularization effects. When
small enough that (GP < 1/2, we expect to observe grokking-like behavior. Specifically, for some threshold < (e.g., = 0.01): let lar such that [V II| x =ll F[oo)( Now let EWD denote the time when weight decay kicks in: i.. when the magnitude of the unregularised gradient becomes comparable to the magnitude of the weight decay term: || V (W(ww)) ]ll Rs X|] AlIWtwa.ll Since at time ŁWD, the solution is close to the gradient flow solution GP
consider Weww Rs ||/ The condition for grokking to occur (i.c., plateau in test loss followed by an improvement of the test loss) is thus that fGF & twD. Translating this condition in terms of gradients, obtain: [V )l >> | F IlI. which, using the approximations above, implies: <l| V([ || > 2l|2 Simplifying further (absorbing into
constant), we have the practical guideline: X < [V Hence, grokking occurs when the weight
[[
decay parameter is sufficiently small compared the ratio between the initial gradient magnitude and the norm of the unregularised gradient flow solution.
37


Reviews:
[{"review_id": "WpdRDbzk3l", "reviewer": "Reviewer_Bgxe", "summary": "The paper lays down purely theoretical justification for grokking phenomenon, a two-phase learning behavior where generalization occurs long after training loss reaches zero. With assumptions on stability, smoothness, and non-pathological behavior of the underlying models, the authors has built a purely theoretical justification of two-phase learning behavior under small l2 weight decay. The paper also supplies small-scale experiments to demonstrate the theoretically discussed behaviors.", "questions": "These are the questions that I do not think to be appropriate to be classified as the weaknesses but do still believe to be important. Below are request for clarifications of my possible misunderstandings.\n\n1. I may have lost but is t≈1/λ transition repeatedly appearing throughout the manuscript justified? Is this transition boundary sharp or gradual in practice?\n2. How small the weight decay should be in order to be applied by this theory? Do the authors have a clue when does the grokking start to not happen as we increase the weight decay OR decrease the initial weights?\n\nThe next question is from my pure curiosity and therefore the answer will not affecting the overall scores.\n\n1. In the Line 347-349 and the referred Appendices C and D.2, the authors have argued that grokking may promote sparse estimators. From my point of view, grokking seems to “correct” the bad initialization at the beginning of the training. As demonstrated in OmniGrok, sufficiently small weight initializations can quickly set the model to the so-called “rich region”, and **de-grok** the model. This implies that we have to be careful about the features that seem to be **unlocked** by the grokking phenomenon. Is this the “grokking” phenomenon that “unlocks” certain features of the model that would not be possible without “large weight initialization & grokking”, or is it just the problem-intrinsic feature that can also be achieved by **de-grokked** models?", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "## Strengths\n\n1. The paper presents (almost) model-agnostic explanation of grokking, achieving generalizability and bridges many of the empirically found grokking examples on diverse set of architectures and data (e.g., OmniGrok).\n2. The assumptions underlying the theory is reasonable, and embraces rich set of models.\n3. The paper is well-formatted and pleasurable to read.\n\n## Weaknesses\n\nMajority of my concerns are inquiries rather than criticisms. But here are some points the authors may consider in order to improve the presentation and connection to the practical branch of works.\n\n1. It would be nicer to have a list of real-world functions/models/losses/training scenarios that the underlying assumptions (Assumption 1, 2) hold. It is roughly argued for generality of the assumption throughout page 4. I believe, however, listing concrete examples, e.g., MLP with Lipschitz-continuous activations like smooth rectifying or sigmoidal nonlinearities, sigmoid after quadratics like self-attention layer, etc., will greatly enhance the impact of the theoretical findings, and help the follow-up works to find out where should they start working on.\n2. I believe it will be much better if there is another list that those assumptions do not hold. For example, Line 129-131 and line 328-330 were graceful, but a dedicated list of exclusion maybe in the separate limitation section would be more elegant.\n3. The first report on grokking phenomenon was on Transformer network, which contains not only residual MLP but also residual multi-head attention. So it will be nicer to have this application explained by the proposed theory, too.\n\nOverall, I believe the paper is already in a good shape regarding the theoretical novelty. However, I have several concerns related to presentation and demonstration. I will raise my score if those issues are resolved. Please also refer to the questions section.", "was_revised": true, "final_justification": "I have read through all the rebuttal and other reviewer's comments. The authors have provided a rich set of explanations that has resolved many of my initial concerns. However, I believe that without further closing the gap between the Transformer architecture and the suggested theory in an explicit way (e.g., establishing Assumption 2 as the authors has suggested), the practical value of this work remains not significant enough. Therefore, I will keep my score.", "rebuttal": "We thank the reviewer for their valuable feedback and will carefully address all the remarks in the revised version.\n\n## When do Assumptions 1, 2 hold?\n\nIn short: our assumptions are quite generic. They hold for most basic neural network architectures with smooth activations, up to some important caveats (in general, they might only hold *locally*, and after excluding certain degenerate configurations).\n\n**Assumption 1** is pretty mild: it holds for all architectures that use differentiable activations (i.e. that do not use ReLU). Although ReLU is a very common activation, many architectures now use differentiable variants of ReLU such as GeLU and thus satisfy Assumption 1. The ‘restrictive’ part of Assumption 1 is that it assumes that $w_{GF}$ is bounded: this does not hold when training with cross entropy loss (or other exponential tail classification losses) on separable datasets. Apart from this case  (i.e. in regression tasks or when the convergence point does not interpolate the data), Assumption 1 holds.\n\n**Assumption 2** (Morse-Bott property) is the more restrictive one. It requires the Hessian matrix $\\nabla^2 F(w)$ to have a constant rank on $\\mathcal{M}$. In general, this is not true globally, as $\\mathcal{M}$ can have degenerate points. However, in most cases it holds *locally* around *almost every point* of $\\mathcal{M}$. For example:\n- for the \"diagonal neural network\" toy model, it holds if we restrict to points with nonzero coordinates,\n- for matrix completion, it does if we exclude matrices with degenerate rank,\n- it also holds *locally almost everywhere* for wide MLPs, ConvNets and ResNets with smooth activations; see for instance Liu, Zhu and Belkin (2021)\n\nIn our work, we preferred to assume that the Morse-Bott property holds globally to derive simpler results and analysis. Similar global assumptions were considered by the works of Li et al (2021), Shalova et al (2024), Fatkullin et al (2010). Although it is not exactly true in practice, we believe it still allows to derive meaningful analyses.\nNote that we could extend our results to the local case, provided that the iterates remain in the non-degenerate region. \n\nWe will add more discussion about these assumptions in the revised version, using the extra page allowed for the camera ready version.\n\n## Grokking with attention explained by the proposed theory\n\nWe agree that applying our theoretical framework to analyze grokking in Transformers is a promising and worthwhile direction.\n\nAs mentioned above, our assumptions are quite generic. They can be shown to hold for attention-based architectures, although establishing Assumption 2 rigorously requires extra work. \n\nHowever, to show that our results lead to grokking in this case, we need to understand the generalization properties of minimal-norm solutions for Transformers.\nIn Section 5, we focus on simple neural network architectures, for which there already exists a rich body of work characterizing implicit bias and minimal-norm interpolators. \nExtending this analysis to attention-based models is a highly non-trivial task that calls for a study of its own.\n\n## Is the transition at $t=\\frac{1}{\\lambda}$ theoretically justified? Is the transition sharp or gradual?\n\nYes, our work gives a precise statement showing that the transitions happens at $t=\\frac{1}{\\lambda}$.\nThis follows from the change of timescale for $\\tilde{w}$, as defined on line 208: we will make this clearer in the revised version.\n \nIndeed, Proposition 2 implies that at time $\\varepsilon/\\lambda$, the regularized flow remains close to the limit of the unregularized flow, which interpolates but does not generalize well in the high-norm initialization regime. In contrast, at time $\\Theta(1/\\lambda)$, the flow approaches the limit of the Riemannian gradient flow on the interpolation manifold. This point has smaller norm and, as a result, is expected to generalize better.\n\nAs explained in our answer to Reviewer CjUp, this transition is gradual (both theoretically and empirically), but appears sharp when time is plotted in log-scale. Though mathematically incorrect (in the sense that the transition is not “immediate” even in log scale), people tend to say it looks sharp. We claim this is mostly because relatively to the long plateau before the transition, the transition indeed looks “rapid”, i.e. “sharp”.  \n\n## How small should the weight decay be to observe grokking?\n\nIndeed, our results are derived in the asymptotic regime $\\lambda \\to 0$, since this allows for a tractable analysis. Extending the theory to a fixed $\\lambda > 0$ is considerably harder. That said, we can offer some heuristic analysis regarding how small $\\lambda$ needs to be for grokking to emerge. \n\nGrokking depends on a clear separation between two phases: an initial phase where the training loss rapidly converges to zero, and a second, slower phase driven by weight decay, during which test performance improves. For grokking to be observable, the gradient flow should approach a near-stationary point before weight decay begins to significantly influence the dynamics.\n\nTo formalize this intuition, we can define two characteristic times: $t_{GF}$​, the convergence time of the unregularized gradient flow, measured as the time at which the gradient norm substantially decreases relative to its initial value; and a second timescale of order $1/\\lambda$, associated with the onset of regularization effects. When $\\lambda$ is small enough that $t_{GF}\\ll 1/\\lambda$, we expect to observe grokking-like behavior. Specifically, for some threshold $\\varepsilon \\ll 1$ (e.g., $\\varepsilon=0.01$): $ \\Vert \\nabla F(w_{t_{GF}})\\Vert  \\approx \\varepsilon \\Vert \\nabla F(w_0)\\Vert $. Now let $t_{WD}$ denote the time when weight decay kicks in: i.e. when the magnitude of the unregularised gradient becomes comparable to the magnitude of the weight decay term: $\\Vert \\nabla F(w_{t_{WD}})\\Vert  \\approx \\lambda \\Vert w_{t_{WD}}\\Vert $. Since at time $t_{WD}$, the solution is close to the gradient flow solution $w^{GF}$, we can consider $\\Vert w_{t_{WD}}\\Vert  \\approx \\Vert w^{GF}\\Vert $. The condition for grokking to occur (i.e., a plateau in test loss followed by improvement) is thus that $t_{GF} \\ll t_{WD}$. Translating this condition in terms of gradients, we obtain:\n$\\Vert \\nabla F(w_{t_{GF}})\\Vert \\gg \\Vert \\nabla F(w_{t_{WD}})\\Vert $, which, using the approximations above, implies:\n$  \\varepsilon \\Vert \\nabla F(w_0)\\Vert \\gg \\lambda \\Vert w^{GF}\\Vert $. Simplifying further (absorbing $\\varepsilon$ into a constant), we have the practical guideline: $\\lambda \\ll \\frac{\\Vert \\nabla F(w_0)\\Vert }{\\Vert w^{GF}\\Vert }$. Hence, grokking occurs when weight decay is sufficiently small relative to ratio of the initial gradient magnitude over the norm of the unregularised gradient flow solution.\n\n**Empirically**, we observed the grokking effect in our experiments for values of $\\lambda$ up to $10^{-2}$. For such relatively large values, the initial plateau becomes noticeably shorter. As $\\lambda$ increases further, the two phases begin to slightly overlap, making the separation between them less distinct (but still present).\n\n## Question on feature unlocking and de-grokking\n\n> Is this the “grokking” phenomenon that “unlocks” certain features of the model that would not be possible without “large weight initialization & grokking”, or is it just the problem-intrinsic feature that can also be achieved by de-grokked models?\n\nOur interpretation of the grokking phenomenon is as follows: when initialized with large weights and trained without weight decay, the model does not necessarily learn meaningful features in the sense of prioritizing the relevant ones. Instead, its output depends on a broad set of features, including many that are not informative. In contrast, when regularization is introduced (or in the so-called \"rich\" regime), we do not believe that the model *“unlocks”* entirely new features that were absent in the earlier regime. Rather, it learns to concentrate weight primarily on the relevant features–those that contribute meaningfully to generalization–while down-weighting or discarding the irrelevant ones. \nThis contrasts with the unregularized setting (with large initialization), where the model tends to distribute weight more uniformly across both relevant and irrelevant features, with the latter possibly dominating due to their greater frequency.\n\nIn this context, the notion of \"sparsity\" refers to the model ultimately placing significant weight on only a small subset of useful features. Prior to grokking, by contrast, the model assigns non-negligible weights to a broader set of features, many of which are not essential for generalization."}, {"review_id": "oMmsg7Mgwn", "reviewer": "Reviewer_CjUp", "summary": "This paper studies the gradient flow with weight decay in the limit as λ→0. The authors characterize its two-stage behavior: in the first stage it tracks the unregularized gradient flow until it reaches the manifold, and in the second stage it follows the Riemannian flow on that manifold. They use this theoretical framework to explain the phenomenon known as grokking. They verify their findings through synthetic experiments.", "questions": "1. How do the authors explain the sudden drop in validation loss observed in grokking, rather than the slow Riemannian flow tracking predicted by their theory?  \n2. In Section 4.2, the definition of the limit flow appears to follow Li et al. (2021) rather than being original, yet it is not properly cited.  \n3. Could the authors provide experiments on real-world datasets? Even small-scale ones would be fine.", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths:**  \n- The paper tackles an interesting problem and provides a new theoretical perspective on the grokking phenomenon.  \n- The writing is fluent, with thorough coverage of related work and background, making it easy to follow.  \n- The simulation experiments convincingly support the authors’ claims.  \n\n**Weaknesses:**  \n1. The main contribution is not prominent. To my knowledge, similar two-stage gradient flow analyses have been widely explored; although the authors’ focus on the weight-decay variant is novel, their theory relies on the λ→0 limit, which substantially weakens the findings. I question the practical relevance of this assumption. Even theoretically, if λ is a nonzero constant, the gradient flow in the first stage may not converge to the manifold at all. Additionally, the proof of the main result (Proposition 2) heavily relies on the definition and properties of the limit flow from prior work, offering limited technical innovation.  \n2. The theory does not fully account for grokking. Under the authors’ framework, the second stage should slowly track the Riemannian flow—implying a plateau and gradual change in validation loss—whereas grokking typically shows a long plateau followed by a sudden drop in validation loss.  \n3. All experiments are synthetic. Experiments on real-world language models are essential.", "was_revised": true, "final_justification": "The author provided a good rebuttal, clearly explaining the plateau period of grokking and providing additional background by connecting it with earlier works. While I still consider the technical novelty and practical significance to be limited, the clarifications have addressed several of my concerns, and based on these considerations, I have raised my rating from 2 to 3.", "rebuttal": "We thank the reviewer for their valuable feedback and will carefully address all the remarks in the revised version.\n\n## 1. How do you explain the sudden drop in validation loss observed in grokking? \n\nWe would like to clarify that the observed drop is not actually sudden, even in practice. It may appear so because, as in most grokking papers, results are plotted in log-time scale. In such plots, the \"drop\" phase occupies a much smaller portion of the x-axis than the preceding plateau, which can create the impression of a sudden change. However, when viewed on a linear time scale, the drop is more gradual. This is not unique to our experiments but is common across many grokking studies (e.g., Power et al., 2022; Liu et al., 2022c).\n\nFurthermore, Proposition 2 in our paper gives a theoretical explanation for this. It predicts that the second phase of learning occurs within a time interval $[\\varepsilon/\\lambda, M/\\lambda]$, where $\\varepsilon$ is small and $M$ is large, both independent of $\\lambda$. As a result, on a log scale, the drop appears within an interval of length $\\ln⁡(M/\\varepsilon)$, while the plateau spans $[\\Theta(1),\\Theta(1)/\\lambda]$, corresponding to a length of $\\ln(1/\\lambda)$ on the same scale. This explains why the drop looks sudden relative to the plateau in log-scale plots. However, in linear time, the drop spans a length of a similar order of magnitude than the preceding plateau, giving a different visual impression.\n\n## 2. Relation to Li et al (2021)\n\nWe acknowledge that Li et al. (2021) revived the “two-stage” gradient flow analysis, which is the foundation of our proof of Proposition 2. We do not claim novelty of this analytical framework. In fact, this approach dates back to Katzenberger (1990) and was further developed in more modern contexts by Fatkullin et al. (2010). The limit flow $\\Phi$ is a standard object in continuous-time optimization and already appears in Katzenberger’s work. We will clarify these connections more explicitly in the revision. In particular, we will emphasize how Li et al. (2021) recently highlighted the power of this technique.\n\nAt a high level, our contribution differs from Li et al. (2021) in that we explain a slow drift induced by a **deterministic regularization**, whereas their analysis focuses on stochastic effects. Additionally, we relate this drift directly to the **grokking phenomenon**, which to our knowledge has not been previously addressed in this context. From a technical standpoint, our setting is simpler and allows for a more detailed analysis:\n1. Rather than invoking Katzenberger (1990) as a black-box result, we provide a self-contained and accessible proof that builds on Falconer (1983), whose ideas also underlie the works of Katzenberger and Li et al.\n2. Moreover, while prior analyses typically assume initialization close to the interpolation manifold, our approach handles **general initialization**. We rigorously analyze the first phase, where the flow approaches the manifold, and then establish the transition to the second phase (drift along the manifold). Dealing with this transition is non-trivial and, as we mention following Proposition 2, is the main technical novelty in this proof.\n\n## 3. Experiments on real-world datasets.\n\nMany empirical studies have already demonstrated the grokking phenomenon on real-world datasets. We reference several of these in the paragraph “Grokking in experimental works” in Section 2, and we do not believe that reproducing similar experiments would yield additional insights. \n\nOur objective is not to contribute further real-world evidence, but rather to validate the coherence between our theoretical results and controlled toy simulations. These simplified settings allow us to precisely monitor the relevant quantities and illustrate our theoretical findings. \n\nNaturally, we would be happy to include any additional references of interest if you have suggestions.\n\n## 4. The theory relies on the $\\lambda \\to 0$ limit.\n\nOur work indeed focuses on the $\\lambda \\to 0$ limit, as it allows for a simple and tractable analysis. Providing clean results with a constant $\\lambda > 0$ would require stronger assumptions and likely lead to much more complex proofs.\n\nRegarding the statement that\n\n> even theoretically, if  $\\lambda$ is a nonzero constant, the gradient flow in the first stage may not converge to the manifold at all, \n\nalthough this is true, we can still show that the flow reaches a point that is *close to the manifold*, and upper bound the distance as a function of $\\lambda$.\n\nIn particular, Equation (9) in the Appendix (proof of Proposition 1) provides an upper bound on the deviation between the regularized and unregularized flows. This allows us to bound the distance of the unregularized flow to the manifold for finite $\\lambda$ and a carefully chosen time $t$ (as detailed in the proof of Lemma 2). While we acknowledge that this bound may not be tight enough for practical use, it does offer some theoretical control in the finite $\\lambda$ setting.\n\nAlso note that in all our synthetic experiments, which are run for non-zero but small values of $\\lambda$, the training loss goes down to *nearly* zero, indicating that the iterates *nearly converge* towards the manifold. \n\nBesides, although our results are asymptotic, they appear to match what we observe empirically for fixed, small $\\lambda$.\n\nWe hope our response clarifies the points raised and appreciate the reviewer’s consideration of these explanations."}, {"review_id": "D4TOX5UMrp", "reviewer": "Reviewer_W5fb", "summary": "This paper gives account of grokking as a dynamical phenomenon that occurs over two timescales: when training with very small weight decay $\\lambda$, gradient flow first (a) converges to very close to whatever solution it was going to anyways, then (b) slowly drifts towards a local optimum of the $\\ell_2$ weight penalty on the manifold of loss minimizers. They prove this for fairly general losses; nothing's specific to deep learning in their analysis.", "questions": "I suppose my main questions are conceptual and mostly listed above. The big one (which isn't easy, and which one could and should ask of most papers) is: why does this matter, really? Why does grokking matter, and who does this paper really help? I think a really convincing answer to this actually would get me to increase my score, but that'd be difficult. I'd be curious about the answer regardless of whether or not I'm convinced.", "limitations": "yep", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Some merits of the paper include:\n* it's clearly written and easy to follow.\n* its conclusion is intuitive (this is emphatically a plus, not a minus), and this makes it easy to skip around and get the gist. the various steps link together in an intuitive way.\n* while I didn't check the math, the authors seem to cross their t's and dot their i's; the notation, presentation, etc. seem both mathematically complete and pretty readable, which isn't easy to do.\n\nAll considered, to me, this reads as a cogent and well-presented story, and it's nice to have a formal reference for this intuitive idea. Seems to clear a reasonable bar for quality and usefulness, so I recommend acceptance.\n\n...but ya can't have it all, so here are some critiques, which I share with the aim of improving the paper and explaining why my score isn't higher:\n* This doesn't seem all that conceptually novel. *That's okay* -- too many papers try way too hard to be conceptually novel and force it -- but I'm just noting it. I've heard this qualitative explanation for grokking before, and the story of this paper was basically what I figured was meant: the weight decay kicks in on a longer timescale, so you get slow movement after naive convergence, and this can slide you over to a better solution. It's nice to see this formalized, especially in deep-learning-independent language, but I'm not sure if there supposed to be any surprises or new insights you get from this analysis.\n* I'm not convinced grokking itself is all that interesting or important. I don't count this against the paper for purposes of this review, but it *is* a broader public conversation that should happen somewhere. Why, again, do we think that grokking matters? If anything, the practical utility of this paper seems to be that, having an explanation for grokking, we can stop studying it so intensely, be satisfied that it's just a result of timescale separation, and move on to other things, which we could have done without the formal description of the phenomenon.\n* If there were a mystery of grokking, seems like it'd be in the test performance: why, with unregularized GF, we'd converge to a spurious solution, but then converge to the \"right\" solution with regularization. (My skeptical lean here reflects the fact that this is a much bigger Q and reaches well beyond the phenomenon of grokking.) From this perspective, seems like we've formalized the easy part, but the deeper mystery, except in special cases, is still open.\n\n\nI hope the authors take the fact that my criticisms are fairly high-level to be praise of the paper: I didn't have to spend time critiquing \"implementation details.\"", "was_revised": true, "final_justification": "This is a high-quality paper. It solidifies a previously-vague idea. My score isn't higher mostly because I doubt the importance of grokking, and the main idea here is one I'd already heard several times.", "rebuttal": "We thank the reviewer for their valuable feedback and will carefully address all the remarks in the revised version. These questions/remarks are very high level and we share most of these concerns/questionings. We give below our “point of view” on these matters, more than fully supported answers. \n\n## Why does grokking matter?\n\nIn that direction, a recent position paper at ICML 2025 [1] argued that grokking may not be a particularly central phenomenon and suggested that the research community should be cautious about devoting excessive attention to it. However, even within that paper, the authors emphasize that studying such settings can contribute to a broader and deeper understanding of training dynamics in deep learning. We fully agree with this perspective and believe our work advances this broader goal. Our theoretical framework is general enough to capture meaningful and widely applicable behaviors in neural network training.\nWhile our main results are motivated by the grokking phenomenon, they also offer insights beyond it. Specifically:\n1. Our analysis provides a general characterization of training dynamics in the small-regularization regime, which can be useful for understanding behavior in a variety of settings—not limited to grokking.\n2. Although grokking is often discussed in the context of weight decay, our approach can easily be extended to other forms of regularization. In particular, we believe that some empirical observations reported in *“The effect of SAM for deep networks at different stages of training”* [2] for training with Sharpness-Aware Minimization (SAM), may also be interpreted through the lens of grokking, albeit driven by SAM-style regularization rather than standard weight decay.\n\n## This doesn’t seem conceptually novel\n\nWe agree that the intuition that “grokking is due to a first interpolating phase, followed by a second phase where regularization kicks in” has been suggested before, for instance in the OmniGrok paper. Since this idea is natural and makes sense, we were actually surprised not to find any formal mathematical treatment of this intuition in the literature. Our work aims to fill this gap. Moreover, while it is nice to say that “regularization kicks in,” as mathematicians, we appreciate having a more precise description. This is what we provide here, showing that this second phase exactly corresponds to a Riemannian gradient flow for the $\\ell_2$​-norm on the interpolation manifold.\n\nAdditionally, we felt that this interpretation of grokking was not universally understood, and that the role of weight decay remained broadly unclear, as raised in our “role of weight decay” paragraph.\n\n\n## we've formalized the easy part, but the deeper mystery, except in special cases, is still open.\n\nOf course, the question of generalization in deep learning is a longstanding and largely unresolved problem that has been extensively studied over many years. In this work, we build upon the commonly held belief that smaller norm weights tend to correlate with better generalization, while acknowledging that this relationship is far from definitive in general settings. Our goal is not to solve this fundamental question, but rather to leverage existing insights on generalization to provide an interpretation of grokking that complements our optimization analysis.\n\n-----------\n# References\n\n[1] Jeffares, A., & van der Schaar, M. *Position: Not All Explanations for Deep Learning Phenomena Are Equally Valuable*. In Forty-second International Conference on Machine Learning Position Paper Track.\n\n[2] Andriushchenko, Maksym, and Nicolas Flammarion. \"*Towards understanding sharpness-aware minimization*.\" International conference on machine learning. PMLR, 2022."}, {"review_id": "aDQ9h8HAFo", "reviewer": "Reviewer_U5vJ", "summary": "This paper provides a theoretical framework explaining how grokking can arise as a result of weight decay. This framework suggests that gradient flow (GF) on loss $F$ with small weight decay has two phases: in the first phase the trajectory quickly converges to a manifold of stationary points of $F$, and in the second phase the dynamics follows the Riemannian GF on this manifold, minimizing the $\\ell_2$-norm of the parameters and slowly converging to an interpolant with a much lower norm, which is believed to generalize better.", "questions": "1. In Figure 3, I could not find the loss curves without weight decay. Did they get lost?\n2. Could the authors clarify a bit the dependence on initialization scale, to provide some connection with previous work that focuses on NTK vs. rich regimes?", "limitations": "Yes, the limitations are discussed openly.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\nIn terms of quality:\n* The framework is very general optimization (in particular, no structure is assumed on the parameters). The assumptions seem reasonable for a theoretical treatment. \n* The treatment is (mathematically) rigorous, the claims are made both informally and formally, and proven in the appendix. \n* The theoretical results are supported by simple examples and experiments, covering the standard toy models (quite convincingly).\n* Limitations are discussed very openly. No overclaiming noticed. \n\nIn terms of clarity, the paper is very well-written. The main theoretical results are split into four propositions and explained well. Toy examples are also clearly reported. There is a nicely written review of related work, with connections well-made (rather than just a list of works).\n\nIn terms of significance, this is an important perspective on the reason for grokking, because of its generality. \n\nIn terms of originality, this perspective provides significant improvement over the closest works (such as Lyu et al. 2023, Liu et al. 2022c). \n\n## Weaknesses\n\nIn terms of significance:\n1. This is gradient flow, where weight decay is synonymous with $\\ell_2$ regularization. Transformers are usually trained by smoothed versions of signGD, where its role can be different (for example, AdamW converges to KKT points of $\\ell_\\infty$-norm constrained optimization https://arxiv.org/pdf/2404.04454).\n2. As authors point out, grokking has been observed without weight decay, whereas the explanation in this article is completely based on weight decay.\n3. Relatedly, classification losses are not handled (also a limitation acknowledged by the authors in the conclusion). \n\nAs a minor point, the proof of Proposition 3 in Appendix A.4 got interrupted mid-sentence.", "was_revised": true, "final_justification": "The authors have agreed to fix the small typos and fix the presentation issues indicated by myself and other reviewers. I believe the limitations are well discussed / clear and the scope of the work is generic enough (although it does not encompass everything and does not completely explain grokking). I would treat this as mainly a theoretical contribution, so I would not hold lack of experiments with real-world language models too strongly against this paper. Overall, after reading the other discussions, I keep my score.", "rebuttal": "We thank the reviewer for their valuable feedback and will carefully address all the remarks in the revised version.\n\n## Loss curves without weight decay?\n\nFor readability, we omitted the loss curves without weight decay from the figure. Overall, the final solution without weight decay closely matches the (regularized) state at time $t_2$​, and the corresponding loss curves align with those shown up to $t_2$​ before remaining essentially flat. We would be happy to include these curves in the revised version if the reviewer believes it would improve the presentation.\n\n## Clarify dependence on initialization scale\n\nUsing the additional page allowed for the camera-ready version, we will include a more detailed discussion on the influence of initialization scale. The initialization scale affects the behavior of the unregularized flow and consequently impacts the first phase of training under regularization. This scale plays a crucial role in the training dynamics, as extensively studied in the implicit bias literature. Although a complete understanding remains elusive, the common consensus—supported by many studies—is that small initialization scales correspond to the rich regime, where implicit bias guides the model toward interpolating solutions with smaller weight norms, which generally leads to better generalization. In contrast, large initialization scales lead to the NTK or lazy regime, where features remain nearly fixed during training. This regime behaves similarly to random feature models and typically results in interpolators with poorer generalization performance.\n\nIn our setting, this means that when starting from a small initialization scale, the point $\\Phi( w(0) )$ reached after the first phase already has a small norm--potentially corresponding to a KKT point of Equation (7)--and thus no second “grokking” phase occurs (i.e., no significant change happens afterward). Conversely, with a large initialization scale, $\\Phi(w(0))$ has a large norm, leading to substantial movement during the second grokking phase.\n\nWe will include additional experiments in the appendix of the revised version to illustrate the impact of initialization scale on the grokking phenomenon."}]

Author rebuttal:

Meta-review: This paper gives a general, optimization-centric account of grokking by proving that gradient flow with small L2 weight decay exhibits two phases: a fast approach to an interpolation manifold, followed by a slow Riemannian gradient flow that minimizes the L2 norm along that manifold. Commendations mainly include broad, mostly model-agnostic theoretical framework with reasonable assumptions (Bgxe, U5vJ), mathematically careful and rigorous (W5fb, U5vJ), empirical support that matches the theory (CjUp, U5vJ, Bgxe), and useful perspective/advance relative to prior explanations of grokking (CjUp, U5vJ). Main concerns include practical relevance of the lambda -> 0 regime and limited novelty beyond prior two-stage analyses (CjUp, W5fb), gap to real-world practice: experiments limited to synthetic data; unclear coverage of transformers/attention (CjUp, Bgxe), and explaining the observed "sudden" drop vs predicted gradual drift (CjUp, Bgxe). Overall, I think this is a nice paper deserving to be accepted.
Decision: Accept (poster)