Grokking at the Edge of Linear Separability
Authors: Alon Beck
Noam Itzhak Levi
Yohai Bar-Sinai
Venue: ICML 2025 poster (2025)
Keywords: Grokking
Logistic Regression
Interpolation Threshold
Memorization and Learning
Primary area: theory->learning_theory
TL;DR: We demonstrate that simple binary logistic regression can exhibits grokking when trained on nearly linearly separable data, initially overfitting before transitioning to a generalizing solution.

Abstract: We investigate the phenomenon of grokking -- delayed generalization accompanied by non-monotonic test loss behavior -- in a simple binary logistic classification task, for which "memorizing" and "generalizing" solutions can be strictly defined. 
Surprisingly, we find that grokking arises naturally even in this minimal model when the parameters of the problem are close to a critical point, and provide both empirical and analytical insights into its mechanism.
Concretely, by appealing to the implicit bias of gradient descent, we show that logistic regression can exhibit grokking when the training dataset is nearly linearly separable from the origin and there is strong noise in the perpendicular directions.
The underlying reason is that near the critical point, "flat" directions in the loss landscape with nearly zero gradient cause training dynamics to linger for arbitrarily long times near quasi-stable solutions before eventually reaching the global minimum.
Finally, we highlight similarities between our findings and the recent literature, strengthening the conjecture that grokking generally occurs in proximity to the interpolation threshold, reminiscent of critical phenomena often observed in physical systems.

Grokking at the Edge of Linear Separability
Alon Beck'
Noam Levi Yohai Bar Sinai
Abstract
Large Language Models (LLMs) (OpenAL, 2024; Google, 2023; Zeng et al., 2022; Brown et al., 2020; Chowdhery et al., 2022: Anil et al., 2023), as well as state of the art models for computer vision (Srivastava & Sharma, 2023), defy expectations and are able to generalize with number of parameters far exceeding the so called interpolation threshold (Kaplan et al., 2020; Schaeffer et al., 2023), Interestingly, these models have been shown to exhibit unpredictable behaviors when changing the number of network parameters, not only with respect to generalization, but also in their learning dynamics.
We investigate the phenomenon of grokking delayed generalization accompanied by non- monotonic test loss behavior in simple binary logistic classification task, for which "memorizing" and "generalizing" solutions can be strictly defined. Surprisingly, we find that grokking arises naturally even in this minimal model when the parameters of the problem are close both empirical and analytical insights into its mechanism. Concretely, by appealing to the implicit bias of gradient descent, we show that logistic regression can exhibit grokking when the training dataset nearly linearly separable from the origin and there is strong noise in the perpendicular directions. The underlying reason is that near the critical point, "flat" directions in the loss landscape with nearly zero gradient cause training dynamics to linger for arbitrarily long times near quasi-stable solutions before eventually reaching the global minimum. Finally, we highlight similarities between our findings and the recent literature, strengthening the conjecture that grokking generally occurs in proximity to the interpolation threshold, reminiscent of critical phenomena often observed in physical systems.
critical point, and provide
One such phenomenon, known Grokking, first observed by Power et al. (2022) during the training of a transformer model on modular arithmetic tasks. Grokking occurs when a model initially achieves perfect training accuracy but no generalization (i.e. no better than random predictor), and upon further training, transitions almost perfect generalization. This phenomenon has garnered substantial attention in recent years (Gromov, 2023; Liu et al., 2023; Xu et al., 2023) due to its striking contrast with naive expectations, whereby over-fitting is generally seen as an undesirable property of models that should not generalize with further training, originally dealt with using early stopping (Prechelt, 1996).
In this work, present straightforward model where grokking naturally emerges, allowing us to identify and analyze its fundamental cause as proximity of the system to a critical point. In our simple yet illuminating setting, the asymptotic optimal solution can always be identified, allowing a sharp definition of notions that are typically ambiguous. such "memorization" and "learning".
1. Introduction
Understanding the relationship between the intrinsic properties of data, the training dynamics of neural networks (NNs), and their ability to generalize is crucial to explaining the success of modern machine learning (ML) algorithms. In particular, highly over- peeemeeramttei models based on the transformer architecture (Vaswani et al.. 2023), such as
We study a typical logistic binary classification problem, with the goal of finding linear separator between two Gaussians with distinct labels. We assume that the Gaussians are well separated along the separation axis and contains noise in all perpendicular directions. Extensions of this setup are considered Sec. 5. We mainly focus on the limit of large N, d - oc while keeping the ratio À d/N fixed, although the results can be generalized.
'Raymond and Beverly Sackler School Physics and Astronomy, Tel Aviv University, Tel Aviv 69978, Israel 'école Polytechnique Fédérale de Lausanne (EPFL), Lausanne, Switzerland. Correspondence to: Alon Beck <alonbk2 maabbk2@gmail.commm. com, Noam Levi <noam.levi@opfl.h..
Our main contributions are:
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
1.
We prove, and demonstrate numerically, that grokking may occur in this setting, and is promoted when is close 1/2 and the separation along the separation

---

Grokking at the Edge of Linear Separability
between the Gaussians is small relative to the noise in
2. Related Work
other directions.
Following the discovery of grokking by Power et al. (2022). numerous studies have attempted to elucidate its underlying mechanisms. Liu et al. (2022) showed that when sufficient data determines the structured representation, perfect generalization can be achieved on non-modular addition task. Other works have identified factors contributing to grokking, including pattern learning (Davies et al., 2023), delayed robustness (Humayun et al., 2024), and transitions from memorization to circuit formation (Nanda et al., 2023), and the role of activation sparsity, weight entropy, and circuit complexity in real-world tasks (Golechha, 2024). Others analyzed the trigonometric algorithms learned by networks after grokking (Nanda et al., 2023; Chughtai et al., 2023; Merrill et al., 2023; Gromov, 2023), and demonstrated similar dynamics in sparse parity tasks (Merrill et al., 2023). Additional works proposed "slingshots" (Thilak et al., 2022) or "oscillaaionss" (Notsawo et al., 2023) as explanations for grokking, whereas others focused on the role of regularization (Power et al., 2022; Liu et al., 2023) and of numerical precision (Prieto et al., 2025), which may significantly impact grokking in certain scenarios. We stress that our work requires none of these in order to exhibit or explain grokking.
2.
We show that this happens because A 1/2 is critical point. That is, for X < 1/2 the model will almost surely asymptotically approach perfect generalization accuracy and vanishing loss, while for À > 1/2 the model will almost surely achieve imperfect generalization accuracy and the population loss will diverge at + oc.
3. More fundamentally, we show that generalization depends on whether the training set is linearly separable from the origin (that is, whether the origin is contained in the convex hull of the training set). In the limit N, d + oc, the data almost surely separable from the origin if X > 1/2 and almost surely false otherwise.
4. Morcover, we show that near the threshold value À = 1/2 (or, more generally, when the data is on the verge of being separable), the dynamics may generically track the overfitting solution for arbitrarily long times before transitioning to the optimal generalizing solution. This behavior manifests as a non-monotonic test loss and delayed generalization, and can lead to divergence of the "grokking time".
Recently, a body of works on solvable models which grok in various settings has emerged. Liu et al. (2023); Kumar et al. (2023) and Lyu et al. (2024) have linked grokking to memorization and transitions from lazy to rich dynamics. Zunkovič & Ilievski (2022); Gromov (2023); Doshi et al. (2024) analyzed solvable models exhibiting grokking and related their findings to the formation of latent- space structure, Xu al. (2023) related grokking to benign over-fitting for ReLU networks on XOR data, Rubin et al. (2024) described grokking and Levi et al. (2023) provide the full dynamical solution of grokking in linear regression. Our work attempts to sidestep external probes, and fill the gap between solvable models and representation learning, whereby we can always identify the optimal solutions, while still solving the dynamics of the model. To accomplish this, our work relies on the results of Soudry et al. (2018); Nacson et al. (2019); Ji & Telgarsky (2019), which analyze the late-time dynamics properties of logistic regression for separable and inseparable data under gradient descent (GD).
5. We construct a simple, one-dimensional model which captures the salient aspects of the problem, and explicitly solve the time evolution of the model parameters for several interesting cases.
first-order phase transition,
The main takeaway from this setup is that grokking happens near a critical point, similar to "critical slowing down" in the physics literature. While further study is needed, we conjecture that this applies to other grokking examples, as was demonstrated explicitly (though not necessarily stated in these terms) in Levi al. (2023); Liu et al. (2023); Gromov (2023); Rubin al. (2023; 2024).
The rest of the paper is organized as follows: Sec. 3 presents our main analysis of grokking as a critical phenomenon, beginning with empirical results in Sec. 3.1, studying the possible solutions in Sec. 3.2, and relating it to linear separability in Sec. 3.3. Finally, we bring together the pieces in Sec. 3.4 to explain why grokking occurs near the critical point. Sec. 4 provides tractable effective model which fully captures the grokking dynamics. In Sec. 5, we discuss generalizations of this setup. We conclude in Sec. 6 and discuss future directions. In the appendices we discuss several generalizations of our results and provide some further proofs and derivations.
3. Grokking in Binary Classification
3.1. Model setup and empirical results
Consider a dataset of N training samples ( C Rd+1, We investigate the case where the data consists of two Gaussian distributions with opposite labels, well separated along the axis of separation but with added noise in all
2

---

Grokking at the Edge of Linear Separability
À 0.05
A 0.48
A 0.52
10a
Train
30-
10
Gen
LOss
10
10-
10-3
ting-time prediction
10"7
10-1
1.0
0.9
0.8
Inseparable
Separabble
0.7
0.6
0.6
ancuuncyy Margn preditcil
10
Long-tima prediction
10a
10
10a
10a
101
+
Grikkingg time
-10
10a
10a
101
10a
10a
102
10a
10a
0.0
0.2
0.4
0.6
0.8
1.0
Epochs
Epochs
Epochs
A
Figure 1. Left panels: Gradient descent dynamics for three different values of A d/N. Loss and accuracy over the train and test datasets, and the time evolution of b(t) and [t)())]]. Grokking is significant only when À approaches to 1/2 from below. We can see that for A > 1/2, ||S]| increases indefinitely and generalization is possible (see Eq. (4)). The parameters N 4- 10 ,
17 = 0.01. The direction of S(t 0) was drawn isotropically with [So|| = 0.1 and b(t 0) = 0. The number of test samples is Noest = 10" Right panels: Top: The norm of the limiting value S.x the separable case > 1/2, as function of X. Note that | diverges for X 2. Middle: the accuracy at the end of the training (in blue), and the predicted limiting accuracy (orange), calculated only using the margin the dataset, see Prop. 3.2. Bottom: The Grokking time, defined here the delay between times when Auain/gen surpass a threshold of 0.9. Grokking time and | diverge near À 1/2. Additional details regarding the experiments can be found in App. K.
other directions. WLOG, we take the separation between the Gaussians to be along the first axis. Explicitly, the distribution of the data is T; N(= Σ) where the covariance Σ is diagnoal with Σn1 = O and Σn = aB for i 1, and μ = (n, O, o... is a vector pointing in the direction of the first axis with magnitude μ. We consider binary logistic regression task with linear model without bias.
identity matrix). In this formulation, it is standard treat the first coordinate of the weight vector as the bias term. A more rigorous justification for this equivalence is provided in App. H, where it is also shown that under this mapping the d- dimimensioonal problem the effective learning rate is faster by a factor of μ²2 The benefit of this equivalence that the analysis of the equivalent model is simpler.
To sum up, for simplicity. throughout this paper will analyze the following equivalent d-dimensional problem: Consider N training samples { C Rdd where æ; ~ (0,, I²Ia) and a > the feature standard deviation (mapped to on/N /u the model above). Consider logistic classification, where all input points are assigned the same label, which for concreteness take {y:) -1
Throughout this work, we focus on the regime a & μe, which allows full generalization below the critical point. We explain the reasoning behind this assumption, and discuss cases where it does not hold in See. 5. We also note that neither Gaussianity nor the assumption of unit covariance is essential for our analysis, and only made for simplicity (see App. G).
Linear model: loss and accuracy. The model parameters are weight vector S E E2 and a bias term E R. The output is a scalar f. = f(x)) = S 2; + b. We optimize the empirical cross-entropy loss L(S, b) and measure the
Reduction to d-dimensional problem with bias.
We consider now the limiting case a = 0, in which the first coordinate of each data point is simply +pp. In this case, the gradient flow dynamics can be exactly mapped to d-dimensional problem with bias, in which all points are assigned the same label. The key idea is that, up to scaling factor, the data points take the form a; = (1, where 2; ~ (0. a Ia) is d-dimensional vector (T is the x dl
3

---

Grokking at the Edge of Linear Scparability
empirical accuracy A(S,b) b). given by!
generalization accuracy is given by:
Agen (S,b) = By..V(0.1) (-
(4)
((S
Σ
ST x; +b)
N
6 1 - erf
(v20|(7))
N
A
Σ
Sth T4 -).
N
where erf is the error function.
Proposition 3.1. Perfect generalizution, i.e,, Egen O
where E(fi = log (1 + ) = log (1 + e is the single sample loss and (z)) is the Heaviside function, defined as (2) = 1 if = 2 0 and (=)) - 0f < O.
and Agen + 1, is achieved only if both b + -00 and b/||s|| + -co. That is, b must tend to negative infinity while also being infinitely large compared to ||]]
Optimizer. Throughout the main text we use gradient descent (GD) dynamics. The effects of other optimizers are discussed in App. F. The GD equations training step with learning rate 77
Proof. It is easily seen from Eq. (4) that the condition Agen - requires b/ ||| + -00. b is bounded, then this can only happen for ||||| + O, but this cannot be since then Eq. (3) implies that Lgen is bounded away from zero. Therefore, perfect generalization implies both b - -oo and b/||S|| - -00.
S - S. VL,
b b, b, -
(2)
In this paper we will focus the gradient flow (GF) limit (17 + 0) of these equations.
The bottom panels of Fig. show that b - -oc at late times in all parameter regimes. However, while ||S|] saturates at a constant value for < 1/2, it diverges when i - co for À > 1/2, and does so at a rate comparable to b, leading to sub-optimal generalization lim A(S(I),b()) b(t) <
Numerical results. In Fig. 1, we show numerical results depicting the gradient-descent dynamics of the model across three values of X = d/N. Notably, we observe a significant grokking effect, both in the oommoootonicity of the test loss, and delayed rise in test accuracy, only when À + d 1/2 (there may be some differences between the grokking observed here and other examples in the literature, but they seem to be superficial- see App. B). In the following section, we explain how de can be interpreted as the interpolation threshold in this setting.
Relation to prior results regarding separability. These results are closely related to the framework developed by Soudry et al. (2018), who studied the convergence of binary classification for linearly separable data, and later expanded by Ji & Telgarsky (2019) for inseparable data. In our case, since the model contains a bias term and all labels are the same, the data is always separable by hyperplane "at infinity". To use their framework, we need to work in extended space of dimension d + 1, where extended weight vector = (S, 6) E R=+11 The network solution at the late stages of training can be obtained as a direct corollary of Theorem 3 from Soudry et al. (2018).
3.2. The generalizing and over-fitting solutions
To understand grokking in this setup, we begin by examining the optimal generalizing solution. Since the support of the input distribution is unbounded and all labels are equal, the model must position all points in Rd on the same side of the separating hyperplane, effectively pushing the decision boundary to infinity.
define the
Theorem 1 (Rephrased from Theorem 3 of Soudry et al. (2018) ). In the setting described above, for any smooth monotonically decreasing loss function with an exponential tail, and for small learning rate, GD iterates will converge at the late stages of training to:
To see this rigorously, we derive expressions for the generalization accuracy and loss. Since the data follows a Gaussian distribution, æ; ~ (O, Ia)), the generalization (population) loss is, by definition:
w (t) wSvM log(t) + pp(t),
(5)
Cgen = -a--(0,,= [log exp exp (S (x++(((xx+))))) (s(z+b)]]
(3)
WSVM argmin (8.6) { ||5)]2+ + b² s.I. ST x; b -1].
STa,
-1]
-Eyoo [log (1 + ++(llls)))
Here, WSVM is the solution to the hard margin SVM problem in the extended d + space and p is a residual vector which is bounded all t.
where we used the fact that ST x ~ (0, 32)) Note that Egen depends only on and |||||. Similarly, the
'The labels do not appear explicitly in L since they are identical for all samples.
Connecting this result to the previous discussion, we see
"Note that the SVM solution in the extended d 1 is not the same as the typical formulation of the Support Vector Machine (SVM) with bias in d dimensions, hecause of the different penalty used for the bias term.
4

---

Grokking at the Edge of Linear Separability
Exponential loss, À = 0.48
0-005
A-0.0
1.0
insasarabl
inseperaplc
0.8
-1
0.9
D:52
|S(()
0,6
10
10
5
X-5
A
0.52
0.4
-5
0.7
Separable
Seporable
-6
0.2
No bias
7
0.0.5
(contormal time)
-30
10"
10a
10a
10a
10a
0.0 2.5
5.0
7.5
10
0|5
Figure 2. Evolution of the model parameters. (left) the distribution of S S// where 5 is the final spatial weight vector that was found using GD dynamics for À 0.05, 0.48, 0.52, 0.8. The parameters are identical to those Fig. We can see that for A 0.05, 0.48 the model does not separate the data (because the data is inseparable) while for 0.52, 0.8 it does. The margin is plotted for = 0.8 Middle panel: [t]]l optimized with GD using the exponential loss given in Eq. (6), with and without a bias term. With bias function of the conformal time (Eq. (8). The two curves follow the same path different rates. The inset shows 4||S and Zoge(y] (right) Optimization paths for different À values, shown the b, o S| plane, For inseparable data diverges (3)
lerm, the result is shown
while bounded, while slightly above the limit separability hoth b and ||]] diverge.
indeed that either [b] and/or ||S| must diverge at infinite training times, and the question is now reduced to the directionality of wSVM-
margin (Prop. 3.2.2). For inseparable training sets (A < 5)) the model will always generalize perfectly: lim b(t) = oc and S saturates on a finite value lim S(I) = co (Prop. 3.2.1).
The generalizing solution, which classifies correctly points in d when WSVM (0, -1). (0 being the d- dimensional zero vector) i.e. when it points in the direction of the bias and the separating plane is at infinity. This is exactly the aforementioned condition b > -oo and 16// ||S||| + co. In contrast, over-fitting occurs when the hyperplane far enough from the data to correctly classify all the training samples, but does not go to infinity. In the extended space, this means that WSVM also contains component in the direction of the data, and the model did not correctly learn the data distribution. In what follows, will show that the factor determining whether we observe grokking in this setup is not the regular separability of data points from one another, but rather "separability from the origin" (or, separability with no bias), defined as follows:
However, for X 1 the training set is on the verge of separability, and S diverges (Prop. 3.4).
Consequently, our main result follows: dynamics may take arbitrarily long times reach the generalizing solution. This is the underlying mechanism of grokking in this setting.
3.3. Separability determines whether the model will generalize perfectly or not
Proposition 3.2. The model will reach perfect generalization if and only the data not linearly separable from the origin. In particular:
Definition data-set {x:) H; E Rdx1 linearly separable from the origin iff there exists vector S Rdx1 such that ST x4 for any i.
1.
If the data is not linearly separable from the origin, then lim, b(t) oo while S saturates on a finite value lim S(t = Sox-
If the data linearly separable from the origin, then lim Agen = 1 erf (MMV)) where M is the margin.
In the rest of the paper, we will use "separable" as shorthand for "separable from the origin". We are now ready to present our main claims regarding the grokking phenomenology presented in Fig. We argue that:
prove Prop. 3.2, we first note that due to the "exponential tail" of the cross-entropy loss, at late times the loss is dominated by samples with large model outputs f
The generalization and overfitting at co depend only on whether the training samples (in Rd) are separable (from the origin) (Prop. 3.2).
S x b, for which the cross-entropy loss {(f) of Eq. (2), approaches the exponential loss (e( el (Soudry et al., 2018; Nacson et al., 2019; Ji & Telgarsky, 2019). Specifically, the exponential loss must converge to the same late time dynamics as the cross entropy loss. Therefore, we will consider the exponential loss for which the calculations
For large N, d, the training set is separable if A and inseparable for < (Prop. 3.3).
For separable training sets (A > }). the model will always overfit, and the limiting generalization accuracy directly related the optimal separating
5

---

Grokking at the Edge of Linear Separability
are tractable,
without bias, i.e. the solution
(S, =
1
Σ
N
S"
w.
(6)
S - argmin ||S]|
s.t. S x; S -1]
N 1=1
The norm of s is the inverse of the separation margin M = 1/ ||
In the gradient-floe flow limit, the induced dynamics are
os Öt -N- N
ee Σ ST ab " N
Biri Ni ot
Σ s' x.
Due to convexity, any convex combination of w and w' will also satisfy the constraints, and since they are orthogonal it also has a smaller norm. The combination with the smallest norm is the global optimum. which is easily shown to be proportional to wsvM x (MM2' -1). That is, bet2
Note that both rates are proportional to common time- dependent scalar el10, We can thus define a so-called conformal time?: T(t) = Jo e which is strictly increasing function of I. In terms of r, the time evolution takes the form
both S and b diverge when 1 + oo and lim; TSttl
M7 Plugging the result into Eq. (4) completes the proof. Next, establish the relation between separability and A.
as θT
7 N
M
S" oE e
7 θT = N
Ob
Σ
S
=
Proposition 3.3. For N, d + oc, and X = d/N, the dataset is separable from the origin with probability 1, if > 1/2, and is inseparable if X < 1/2.
The importance of this change of variables is that the dynamics of S()) in terms of the conformal time are identical to those of S(t) in the absence of bias. That is, S()) follows the same path that would be obtained by minimizing = ,o4 EN os'a:, but does so at different the current value of b(t). This is demonstrated in the middle panel of Fig. 2. Since r() diverges for / + oc (see App. C for details), S(t) must follow the same path at long times, as it would have followed without bias. We can now complete the proof:
In other words, Prop. 3.3 states that (almost) any large set of N points in d dimensions are separable (i.e., lie on the same "half" of some hypersphere passing through origin) as long as d > N/2. This is direct corollary of Wendel's theorem (Wendel, 1962) which we prove in App. A.
rate which depends exponentially
3.4. Collecting the pieces: why does grokking happen near À = /
We have established that for À < 1. the model will almost surely generalize perfectly. For infinitely long times, Stt) converges to finite vector Soc. and b(t) diverges. For A 1. the model will almost surely overfit. Intuitively, one should expect that in the vicinity of the critical point d = ) where the two solutions exchange stability, dynamics may become slow. This is because for À >>> 1 1/2 the overfitting solution is stable and for À smaller than but close to 1/2, it is unstable but only slightly so. Therefore, the dynamics may spend arbitrarily long times in the vicinity of the overfitting solution before flowing to the generalizing solution. This delayed generalization. Rigorously, this happens through of the following properties:
Proof of 3.2.1 (inseparable case) Consider the dynamics without bias. In case the data is inseparable, Le of Eq. (6) is unbounded in all directions of S. That is, for any unit vector e E Idd we have lima-toc Le(ave, 0) = oo. Since C. convex, the gradient flow dynamics will lead to a global minimum at a finite point lim (t) = Sx. Recalling the discussion of conformal time above, this is also the limit of the dynamics with bias. From Eq. (5), we know that either [[S())) or [b(L) must diverge, and since (L) approaches a finite value, we conclude that [bl/ ||||| - de for t + co. That is, WSVM = (0, -1) and the model flows towards the generalizing solution.
Proposition 3.4. For - 1T, ||50o - 0c.
Proof of 3.2.2 (separable case) When the training set is separable from the origin, it is easier to examine the optimization problem in Eq. (5) directly. We wish to minimize ||ul|| = ||S| + 12 under the separability constraints. The generalizing solution w, = (0, -1) 1.
That is, when the training set is non-separable, but on the verge of separability. ||S obtains arbitrarily large values. This statement is formally proven in App. D and empirically demonstrated in Fig. 1. It can also be obtained as a corollary of Ji & Telgarsky (2019). An intuitive geometric interpretation is that for a nearly separable set, St) approaches finite limit Sx., but small translation of the data would make the set separable, and correspondingly would make |S())] + oc. Smoothness thus implies S must be large if the set is almost separable.
satisfies all constraints trivially and has ||w,]| However, since the data is separable from the origin, there exists another solution to the constraints, namely w" = (5*,). where 5t is the separating vector in d dimensions
'This is common measure in cosmology and gravitational physics to describe co-moving objects in an expanding shrinking spacetime background (Guth, 1981).
Proposition 3.5. For X < 1 and a large enough, S(t will approach its asymptotic value Sa arbitrarily fast.
6

---

Grokking at the Edge of Linear Separability
0.75
0.5
2
2
2
1el
1.0
1
%
1
1
Bo+ 00+07
76-40
0.8
66407
Sex07
0.6
0
-ugVi
0
0.4
0.2
2 -1
2 2
1
0
1
2
2
-1
0
1
2
8
10-
10
10-
10
10-2
5
S
Figure 3. Simplified model. Three left panels: Illustration of the hard margin SVM problem in I+1 dimensions for the simplified model. Note that - WWSVV as is the point closest to the origin in the intersection of the two shaded regions. When wsv as points along the bias axis (the vertical axis) if and only if 22 2 O. (right) Grokking time, defined as the time it takes for the generalization accuracy to reach 0.95, plotted against a and 2X. see diverges when both À -+ 1/2 and - co, while neither condition suffices alone.
A < 1 1/2 (02 > 0)
A = 1/2 (02 = 0) - log((t)
A 1/2 (02 < 0)
b(t || > 1)
log(t)
M2 log(t)
A) log
log(log(t))
M log(t)
1-2A
I+
Table I. Summary of the different regimes of the simplified model, where x -1,, 11,12 2X, and the margin is M |aaa| = 2A 1.
Proof. To see this, it is useful to define the rescaled variables क; 0:/o, S = oS. Clearly, H; ~ (0, Ia)- The gradient flow equations in terms of the rescaled variables are (we study the exponential loss for simplicity)
4. Insights From a Simplified Model
Our main claim is that the asymptotic dynamics depend only on the separability of the training set, and that grokking occurs at the edge of linear separability. This intuition can be worked out explicitly in a much simpler setting in one dimension. In this case, separability boils down to asking whether the origin is contained between the extremal points min(#;) and max(x;). Therefore, all the phenomenology of the full model described in the previous sections can be captured by training set consisting of only points 31, 22, representing the points with maximal and minimal projections along S For consistency with the problem of Gaussian data, we parameterize this set as
os
01
=
o² 2
N
Ri:
(10)
Öb
n
S"
1
Note that these are identical to the gradient flow equations of the original variables given by Eq. (7), but the dynamics of S are faster by factor of a². Thus, by taking large o, S will approach its asymptotic value So arbitrarily fast, while the dynamics of b will not change. Recalling that o was mapped to OB/ / I in the original two-Gaussians model introduced at the start, we note that a large a is equivalent to introducing large noise in directions perpendicular to the separation axis.
07 = - I2 = o (1 (1-21)
(11)
1 (e Se + S024h))
C(S = so that the scale of x; is o, and they are separable (inseparable) for À < 1/2 (A 1/2). We note that the margin of the dataset from the origin has the same dependence as in the Gaussian model with N points in d dimensions.
We can now understand mechanistically how grokking occurs. For À values close enough 1 from below, the limiting norm ||S5ol| is arbitrarily large (Prop. 3.4). For large enough a, S(t) will grow arbitrarily fast towards S. (Prop. 3.5). Under these conditions, the growth rate of b(t) remains boudned, and the generalization can be delayed for arbitrarily long times. Note that this necessitates both A - 2 and a + co, as is also demonstrated in the right panel of Fig. 3. Interestingly, using adaptive momentum based optimizers like ADAM (Kingma & Ba, 2017), one can see significant grokking even for o = 1, see App. B and App. for more details.
The asymptotic dynamics of this model qualitatively, and sometimes quantitatively, capture the phenomenology of the full problem. The model is fully tractable analytically and the detailed analysis is presented in App. E. We summarize here the main results:
The left panels of Fig. 3 show the geometry of the problem in 1 ++1 dimensions. is easily seen that the optimal SVM solution is s = 0, = -1 if and only the data is not separable, i.e. when the segment 22 2
7

---

Grokking at the Edge of Linear Separability
contains the origin.
only modifying S₁ later times, leading delayed generalization. In Fig. 4 we present numerical simulations supporting this behavior: While σ 0.01 is small enough to closely follow the behavior of σд = 0, for larger values (a. = 0.05, 0. I), we observe that the limiting accuracy is below and the loss begins to diverge at long times.
The limiting value So can be easily found to be Sec Iイーー log (1 2A), for the separable case A 1/2. Indeed, it diverges logarithmically at À 1/2, in agreement with the numerical results of the full model presented in the upper-right panel of Fig. 1. The long time dynamics of [[)())] and b(t) are summarized in Table 1.
Finally, we note that while the precise dynamies depend on the data distribution and labeling scheme, similar results are expected for broad class of linear binary classification problems, long as À is near the critical point. For instance, in App. I, we analyze a generalization of the setup from another perspective, where the constant labels are explicitly modified to be discriminative. The results closely resemble those presented here, see Fig. 16.
The behavior of the loss and accuracy of the simplified model as a function of À is remarkably similar to that of the full model, see Fig.
App. E.
Criticality. We note that this result bears striking resemblance to that of Levi et al. (2023), which employed the MSE loss in linear regression problem, again for N points sampled iid from an isotropic Gaussian distribution. In their setting, the interpolation threshold is at À = 1, in the sense that for À < the model always generalizes asymptoticallyy and never generalizes for À 1. They also found logarithmic divergence of the "grokkiig time" (the time difference between the times it takes for the generalization and training accuracy to reach a certain threshold). It diverges as a function of the distance from criticality as x log (1 vX), which was explained in terms of "critical slowing down" effect, arising from vanishing eigenvalue of the data covariance near criticality. While the two problems quite different, they both display critical behavior near effective interpolation threshold of the corresponding problem. We believe this is not coincidence but rather manifestation of a deeper relation between the behaviors of NNs in the vicinity of critical points.
10a
train,7A00
LOss
test, 6 0
train,
test, 0.01
t 0.05
10
test, a.- 0.05
train, 0 0.1
test, a=== 0.1
102
10'
102
10a
Epochs
1.0
0.9
Accuracy
5. Extensions of the Setup
0.8
In our original setup, assumed σ < μu, which allowed us to reduce the model to a d-dimensional problem with bias term under constant labels. In this section, we explore what happens when this assumption does not hold. We show that while the condition O < μ (or equivalently, constant labels) is essential for full generalization, grokking in the sense of delayed generalization well as non-monotonic test loss behavior, can still be observed provided the system remains near the critical point of linear separability.
0.7
0.5
10a
20:
10a Epochs
10a
0,
0, 0.01
Typically. binary classification models fully generalize only when the number of samples is sufficiently large, i.e., N > However, grokking is observed outside this regime, with d/N z 1/2. Here, the assumption A & μ becomes particularly useful, as it ensures perfect test accuracy for any below the critical point. When this assumption is relaxed, full generalization is no longer guaranteed. Nevertheless, the underlying mechanism discussed in earlier sections remains: near the critical point À = 1/2, the model can initially reduce the loss by adjusting the direction and norm of the perpendicular weights S2,.., S4t1, while
-0.2
0.0
0.2
-0.2
0.0
0.2
O,-
0, P
-0.2
0.0
0.2
-0.2
0.0
0.2
Figure 4. Top: The GD dynamics of the loss and accuracy for the 2-Gaussians model, with different values of a The rest of the parameters are À 0.4, μ 0.2, TB and 77 = 0.1. Bottom: illustration of the scparation for each d.
8

---

Grokking at the Edge of Linear Separability
6. Discussion, conclusions and limitations
2022).
We studied the dynamics of gradient descent a simple setting of logistic classification with strong noise. We have shown that in this setting, grokking occurs near critical point in the asymptotic dynamics. Specifically, at the critical point, which occurs at À 1/2 in our case, the overfitting and generalizing solutions exchange stability. We showed that this non-analytic change in the asymptotic dynamics is the cause for grokking, much like in Rubin et al. (2024); Levi al. (2023); Doshi et al. (2024), and to some extent also Humayun al. (2024), who showed that grokking occurs near phase transition.
Lastly, we did not study the prospect of nonlinear logistic regression, which is closer to deep learning models in the wild. We believe some of our results may be generalized, provided we accept "feature map" description of the model up to the last layer, and consider the SVM solution on the learned features.
Acknowledgments
We thank Hillel Abaroni. Francois Charton, Amit Moscovich, Daniel Soudry and Matthieu Wyart for fruitful discussions. YBS was supported by research grant ISF 1907/22 and Google Gift grant. NI. is supported by the EPFL A14science program.
Intuitively, in the vicinity of the critical point there are "flat directions" in the loss landscape. These directions may cause training to stay in the vicinity of almost- stable solutions for arbitrarily long times periods before eventually converging to the global minimum. In the physics literature, this behavior is known as "critical slowing down" (e.g. (Sethna, 2021)). In the current context, this is the mechanism of delayed generalization, which also explains the non-monotonic evolution of the generalization loss.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, in particular due to the large model sizes considered in this work, but we do not feel there are specific aspects of this work with broader impacts beyond the considerations relevant to all large machine learning models.
While we cannot show it rigorously, we conjecture that grokking is intimately related to such critical points also in different settings. In a few examples, this has been directly demonstrated, (Levi et al., 2023; Rubin al., 2023; 2024; Humayun et al., 2024). We note that other intriguing phenomena, such as the non-monotonic dependence of asymptotic performance on model complexity, a.k.a "double descent", have also been proposed to be related to criticality, e.g. (Schaeffer et al., 2023).
References
Anil, R., Dai, A. M., Firat, O., Johnson, M., Lepikhin, D., Passos, A., Shakeri, S., Taropa, E., Bailey, P., Chen, Z., et al. Palm 2 technical report. arXiv preprint arXiv:2305. 10403, 2023.
If this is indeed the case, then in analogy to the theory of critical phenomena in physics, there might exist "universality classes" that have similar critical behavior, but possibly very different underlying mechanisms (Sethna, 2021). We will address this connection in future work.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
We considered a specific problem of linear binary classification in high dimensions. is natural to ask how our results extend to more complex data, for instance including non-trivial correlations, hierarchical structure, or finite sample space, as in the original observation of grokking in Power et al. (2022); Gromov (2023). While we believe the same analysis can be repeated in these instances, in the sense of (non)linear separability, we leave this to future work. any case, we do not claim that the underlying mechanism of criticality necessarily related to separability.
Limitations:
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.023111 2022.
Chughtai, B., Chan, I.., and Nanda, N. A toy model of universality: Reverse engineering how networks learn group operations, 2023.
Cohen, J. M., Kaur, S., Li, Y., Kolter, J. Z., and Talwalkar, A. Gradient descent on neural networks typically occurs at the edge of stability, 2022.
The analytic were all done in the GF limit, and while our results were verified by experiments with a finite learning rate, it may be interesting to study how large learning rates affect this setup, possibly relating to catapults (Lewkowycz et al., 2020) or the edge of stability mechanism (Cohen et al.,
Cover, T. M. Geometrical and statistical properties of systems of linear incqualities with applications in pattern recognition. IEEE Transactions on Electronic Computers,

---

Grokking at the Edge of Linear Separability
BC-14(3):3263334 1965. doi: 10.1109/GGEC1965 264137.
Liu, Z., Michaud, E. J., and Tegmark, M. Omnigrok: Grokking beyond algorithmic data, 2023.
Davies, X., Langosco. L., and Krueger, D. Unifying grokking and double descent. arXiv preprint, (iXiv:2303.66177 2023.
Lyu. K., Jin, J., Li, Z., Du, S. S., Lee, D., and Hu, W. Dichotomy of early and late phase implicit biases can provably induce grokking, 2024.
Doshi, D., Das, A., He, T., and Gromov, A. To grok or not to grok: Disentangling generalization and memorization on corrupted algorithmic datasets, 2024.
Merrill, W., Tsilivis, N., and Shukla, A. A tale of two circuits: Grokking as competition of sparse and dense subnetworks, 2023.
Golechha, S. Progress measures for grokking real-world tasks, 2024.
Nacson, M. S., Lee, J. D., Gunasekar, S., Savarese, P. H. P., Srebro, N., and Soudry, D. Convergence of gradient descent on separable data, 2019.
Google. Bard chat based ai tool from google (october 2023 version) [large language model). 2023. URL https: //bardd google.com/
Nanda, N., Chan, L., Liberum, T., Smith, J., and Steinhardt, J. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.0522 2023,
Gromov, Grokking modular arithmetic. arXiv prerint, MoXiv:2303.02679 2023.
Notsawo, P. J. T. Zhou, H., Pezeshki, M., Rish, I., and Dumas, G. Predicting grokking long before it happens: A look into the loss landscape of models which grok, 2023.
Guth, A. H. Inflationary universe: A possible solution to the horizon and flatness problems. Phys. Rev. D, 23:347-356, Jan 1981. doi: 10.1103/PhysRew..
OpenAL. Gpt-4 technical report, 2024. URL https:// azxiv.org/ www//bss/ss...ooooooo./.....e.ee abs 08774.
23.347. URL https: //tnk//ink////////////////////./ aps crg/doi/10.. 1103/PhysRevD) 23.347.
Power, A., Burda, Y., Edwards, H., Babuschkin, I. and Misra, V. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:201.00/777, 2022.
Humayun, A. I., Balestriero, R., and Baraniuk, R. Deep networks always grok and here why, 2024.
Ji, 7. and Telgarsky, M. The implicit bias of gradient descent on nonseparable data. In Conference on learning theory, pp. 1772-1798. PMLR, 2019.
Prechelt, L. Early stopping-but when? In Neural Networks, 1996. URL https: /appi.ppppp///// semant escholar. semanticscholar. org //tp//oo::
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models, 2020.
Pricto, L., Barsbey, M., Mediano, P. A. M., and Birdal, T. Grokking the edge of numerical stability, 2025. URL https://ae /ttass///xxxivi wwwerxivv..gg/aee
org/ abs 2001 04697.
Rubin, N., Seroussi, I., and Ringel, Droplets of good
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization, 2017.
representations: Grokking a first order phase transition in two layer networks, 2023.
Kumar, T., Bordelon, B., Gershman, J., and Pehlevan, C. Grokking as the transition from lazy to rich training dynamics, 2023.
Rubin, N., Seroussi, I., and Ringel, Z. Grokking as first order phase transition in two layer networks, 2024.
Schaeffer, R., Khona, M., Robertson, Z.. Boopathy, A..
Levi, N., Beck, A., and Bar-Sinai, Y. Grokking in linear estimators solvable model that groks without understanding, 2023.
Pistunova, K., Rocks, J. W., Fiete, I. R., and Koyejo, O. Double descent demystified: Identifying, interpreting & ablating the sources of deep learning puzzle, 2023.
Lewkowycz, A., Bahri, Y., Dyer, E., Sohl-Dickstein, J., and Gur-Ari, G. The large learning rate phase of deep learning: the catapult mechanism, 2020.
Sethna, J. P. Statistical mechanics: entropy, order parameters, and complexity volume 14. Oxford University Press, USA, 2021.
Liu, Z., Kitouni, O., Nolte, N. S., Michaud, E.,, Tegmark, M., and Williams, M. Towards understanding grokking: An effective theory of representation learning. Advances in Neural Information Processing Systems, 35:34651- 34663, 2022.
Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S., and Srebro, N. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70):1-57, 2018. URL http: / jmlr. prg/paperrs/v19/1e-a html.
10

---

Grokking at the Edge of Linear Separability
Srivastava, S. and Sharma, G. Omnivec: Learning robust representations with cross modal sharing, 2023.
Thilak, V., Littwin, E., Zhai, S., Saremi, O., Paiss, R., and Susskind, The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint arXiv:2206. 048/7, 2022.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,, Gomez, A. N., Kaiser, L.,, and Polosukhin, Attention is all you need, 2023.
Wendel, J. A problem in geometric probability. Mathematica Scandinavica, 11:109-112, 1962. URL http: http://ddmllll.l./ www.or/////o//////
Xu, Z., Wang, Y., Frei, S., Vardi, G., and Hu, W. Benign overfitting and grokking in relu networks for xor cluster data, 2023.
Zeng, A., Liu, X., Du, Z., Wang, Z., Lai, H., Ding, M., Yang, Z., Xu, Y., Zheng, W., Xia, X., et al. Glm-130b. An open bilingual pre-trained model. arXiv preprint 0rXiv:2210.02444 2022.
Zunkovič, B. and Ilievski, E. Grokking phase transitions in learning local rules with gradient descent, 2022.

---

Grokking at the Edge of Linear Separability
Appendix
A. Separability and Wendel's Theorem
Wendel's theorem (Wendel, 1962) states that the probability that N random vectors drawn from a distribution in d dimensions are linearly separable,
N-1
N
p
2
Σ
k
(12)
k=0
In relation to our work, the only assumptions required from the distribution is that
It is symmetric around the origin, i.e. P(æ) = P(-a:), and
The dataset is almost surely in general position.
We note that Eq. (12) is the cumulative probability function of the Binomial distribution, i.e. the probability that the the number of successes is greater than d out of N 1 attempts with success probability 1.2 The central limit theorem states that in the limit of large N, d the binomial distribution approaches Gaussian, and thus the cumulative distribution function approaches the error function. Straightforward manipulations show that for large N, d.
(()) + _ 1 + erl erl
(Va(veL) V2X
)
(13)
N
is seen that for d + oo the transition becomes infinitely sharp as function of lambda and we have
0
lim p(A)
1 X
(14)
1 A >
See also Cover (1965) for further discussion.
B. Relation to canonical examples
In this section, we will discuss the similarities and differences of our work with previous examples of Grokking in the literature, focusing on the seminal work of Power et. al. (Power et al., 2022). We first note that Grokking at Power et. al. is significant when the fraction of the data used for training a = Ntreining/ /N is near a critical value ae, in the sense that the system achieves perfect generalization if and only if > az, as can be seen in Fig. I (center) of their paper. We expect that this non-analytic behavior in the long time limit of training will be the crucial property that underlies grokking. That is, we expect that near such points the dynamics will be slow. We note that Power et. al. analogous to our parameter, defining effective "interpolation threshold" for the modular arithmetic problem.
Secondly, we note that noticeable difference between our work and that of Power et. al. is that in our case, the accuracy shows rise from the start rather than staying at chance level for long time before generalizing. We argue that this is only a superficial discrepancy that depends on the choice of optimizer and fine-tuning of hyperparameters, and that the fundamental mechanism (that grokking occurs near critical points in which solutions exchange stability and dynamics are generically slow) the same.
Indeed, in Fig. we show that our setup is capable of grokking with accuracy staying at chance level (50%) the start, similar to Power al. We achieved this by using À values closer to half ("almost separable") and the Adam optimizer instead of vanilla gradient descent (GD). The fact that this optimizer converges faster on the training data is coincidence: the adaptive learning rate leads to quicker convergence to large values of [S] (the "memorizing solution"), maintaining accuracy at chance level until later stages, before going to large -b (the "generalizing solution").. Notably, Power et al. also used Adam (or AdamW). In conclusion, although Adam can lead to a slightly "cleaner" grokking result, we explored GD because it is easier to derive analytical insights from it while, we believe, not changing the underlying mechanism of
12

---

Grokking at the Edge of Linear Separability
A 0.495
Train Test
Loss
1.0
Train
Accuracy
Test
0.8
0.6
10
0
-5
-10
10
10a
102
103
104
Epochs
Figure Grokking in similar setup the results in the main text with ADAM optimizer (with 31 0.8, ,8 0.9), instead of GD. The parameters are À = d/N 0.495, N 4000 and 1.
grokking. Finally, will also note that the non-monotonicity the test loss also a typical sign Grokking that can be seen in our setup (for example, compare Fig. 4 of Power et with the test loss in Fig. 5).
C. Divergence of the conformal time
In the main text we have defined the "conformal time" r T= e(t)) and saw that as the result the gradient-decent
trajectory of (t) is the same one that minimizes the exponential loss without bias: L = a Σ eSa2, However, if is bounded it might reach a different fixed point. We will show now that indeed must diverge. First, we notice that the loss ST that must be bounded from above: If the points are not separable (that is, there is some 0 such that for any ||S
S choose I. > e for any i), then it must be true since ||S]| is bounded - otherwise the loss would be infinite. If the points separable, then ||S]] might diverge (and will, as discussed in the main text) but some point all of the arguments of the exponent would be negative, the loss would be trivially bounded by 1. Now, using the fact that 82 ØF 3 2b
we have
dz
32 N E S.v. Denoting C(t) C. we see that
Bt
1 08 B a
n
(15)
We note that on the left-hand side we have a positive function (since 0g < 0). In other words, à that
< 1,C, we get
30
+ I nc + nCt
3(0) + Ot
<
(16)
B(t)
3(t)
3(0)
3(0)
so that
1 n't or, (() 1490 This means that fo 3(t) d 14ner which diverges.
F(t)
13

---

Grokking at the Edge of Linear Separability
1.00
0.2
0.95
0.90
0.1
Fraction of ST X; >
0.85
Separable fraction, À 0.5
0.0
0.00
0.05
0.10 1 21
0.15
0.20
10
10
1 21
ST
Figure 6. Left panel: The fraction of positive T8 2i, which goes to constant for À = 1/2. Right panel: The fraction of separable datasets for < 1/2 that were not included in the calculation.
0.2
Avg(S
Margin
0.1
0.0
-0.1
-0.3
-0.2
-0.1
0.0 21
0.1
0.2
1
Figure 7. Numerical investigation of properties of the limiting distribution of S w. as a function of 1 2A (averaged over different
random configurations). In blue, plot the average value of positive the margin for À 1/2 (using SVM).
T Di. for 1/2 (by minimizing L 1 DN s and
s7
D. Proof that S.o diverges for almost separable data
We look at the function
f(S, (xi+)
S-z:
S E IRd
(17)
We will assume n > and that the data is in general position, and that it separable from the origin. Since for every must have S .r 0 for some i, it easy to see that diverges when S grows large in any direction. Since f 0, there exists global minimum at finite S.
A minimum (which is also unique under our assumptions but that's 's not crucial) obeys
0f
2 S-z: 0
(18)
0s
If we divide this expression by f, we get
e S-ze
Pszt
P;
0 05 51,
Σn=1
(19)
Eq. (19) means that the origin is a convex combination of the sample points with weights ps. We found that a necessary condition for the existence of critical point at a finite is that the origin contained the convex hull of the sample
14

---

Grokking at the Edge of Linear Separability
points. This is of course equivalent the condition that the origin is not linearly separable from the sample data.
We want to show that if the data is almost separable, that is, if it is not separable but the origin is close to the boundary of the convex hull, then S must be large. The intuition for this comes from Eq. (19): if the origin is very close to the boundary of the convex hull then some of the p; 's must be very large compared to the others, which can only happen if S is large.
In fact, the origin is exactly on the boundary the convex hull (that is, the data is exactly on the edge separability) if and only if for every representation of the origin as a convex combination of the sample points,
Σ q: = 0,
(20)
the weights 4 are non zero only for le sample points, say Iz..... TR. with S d. and .1.... oy are the vertices of facet of the convex hull. This naturally leads the definition:
Definition We say that the origin is c-close the boundary if there exist l points Fi,---, such that for every representation of the type of Eq. (20), the total weight assigned to I2,.... Ty is at least
k
Σ
2 211
Theorem If the origin is e-close to the boundary of the convex hull of the sample points, then the norm of S = argmin / bounded from below by
1-
S D log
where D = maxiy [x; 2yl is the diameter of the data.
Proof. We divide the points to two groups: A = (x1,.... .at,....k,,.. and B = (Fx+1,...INN) ...... Since the origin e-close, the ratio the weights of the two groups bounded by
ΣiEAPs P:
1-c
(21)
Σien Pi
Consider the convex combination Eq. (19). Using Jensen's inequality, can bound the relative weights of the second group by
Σ S-z. (N- 5-
with
k: x4 B
(22)
B
where & is the average of the points in the second group. Therefore, the ratio is bounded by
EIEAP:
EiEA S...
< EiEA S N - N k.
k SSSD <
k
GSID
(23)
Σien Ps Σien .Siz
(N K)e52
Combining Eq. (21) and Eq. (23) we get
1 1-6 <
k
S
|S| 2 log
1- N - k
(24)
E N
Since 1: < d, we also have N - k 2 nd
Note that the same convexity argument would work also for logisitic loss f -- (( x2).. ((2) = log (1 + e²), or any other monotonic and convex &. this case the only difference is that the log function should be replaced the inverse of
15

---

Grokking at the Edge of Linear Separability
10a
Train Generalization
10
10
Loss
Agen
A
1024
10
0
10
104
1000
2014
0.7
À 0.505
10-1
10%
10
10a
predition
10
10?
10
10"
10
101 107
107
10a
10
10
10
10"3
10"2
10
t
1
2A
Figure 8. Simplified model (left, center) Loss and accuracy for different À and a 5. dotted/solid lines represent the maining/generalizzttonn respectively. (right) Grokking time (the time difference between the time it takes for the training and generalization loss to reach a certain threshold, = 10 in this case) for a = 20. The data is in very good agreement with the prediction. (inset) how grokking time calculated for 10 10-50 and À 10
E. Details of the Simplified Model
E.1. Justification and Relation to the Full Model
We will provide here supplemental results regarding the justification of the simplified model (by numerical comparison to the full model). To obtain the results we average over different random realizations: Assuming the so-called "self averaging" property. we know that the average over large number finite systems should give us the same result as the infinite system (where N, d and the ratio is constant).
In the non-separable case, S.x be found by any optimizer that minimizes the loss L - We note that when getting close to the transition point, for any finite-sized system have probability of getting separable set (even though À 1/2), see Eq. (12). In this case, just ignore the the result: In the right panel of Fig. 6 we present the fraction of realizations that are separable. This will probably introduce some bias into the results which is likely the cause of the fact that the average of positive samples (and similarly, the margin) does not go exactly to zero for À + 1/2 (see Fig. 7). In the || x; Interestingly, does not go to zero but to some positive ST 2 À 1/2.
s"
left panel of Fig. 6 we present the fraction of positive constant, implying that there is singularity in the density of S
E.2. Analytical Predictions
Here, we provide the full analysis of the model presented in Sec. 4, for single point fixed at 41 = -1, and a second point D2 = x = - 2A, where d/N.
The gradient flow equations in conformal time are given by
Øs ?
xe
Sa - E -s ==
17 (1- 21)e 1551-11221111 -
-s-))
(25)
ÒT 2
Ob "
Sz + 55
17
e (11-22)) +e - ---s.
θr 2
2
While there exist analytical solutions for Eq. (25), they do not necessarily provide any intuition, and so we find it better begin by investigating three special representative cases:
1. Z2 (non-separable).
2. x2 (marginally non-separablee.
3. x2 -1 (separable).
16

---

Grokking at the Edge of Linear Separability
For (A); the data is entirely non-separable in one dimension and the conformal time solutions are
S(T) log
(anh + tanh
(c5))))
(26)
(anh (2 tanh -1 (e:50) ) cosh (2 tanh -1 (e:0)
b(T) 6o + log
tanh (1T + 2 tanh
(E cosh (17 + 2 2annh (eSo)
in which case the generalization accuracy reaches I for r + oo, as b(T) grows faster than S(T) with conformal time.
For 22 = 0 (B), the equations in conformal time become:
as 77 - θr 2
bb 11 (c- S + 1)
(27)
Òr
By solving for S and plugging into 8b immediately get
5 log
e 8o
T)
!
log 5050 /) 2/2
!7 So + bn-
(28)
Using the fact that e²
Dr we get that 0r at So 21 + (e 1). Taking the inverse of this, we finally get
4 27
So and taking another integral, we get that e [ - 1+ 2
2
Wo ((e50tha 2 (E
1)) e -00 (650
co + +1]
(29)
where Wo is the Lambert W function. We note that for large t we have T ~ log(t). and therefore b ~ - log (t), ~ log(log(t)), so is interesting to note that the critical point we still have lim, S(t)/b(t) = (i.e., accuracy goes to 1), even though diverges.
Finally, for
(C), the data fully separable in one dimension and the solution in conformal time given by
S(r) So log (1 nyre-S), b(T) = by log (1 nyte" (yre-5)) -0),
(30)
showing that the accuracy is bounded Agen (1 + erf (+2 agreeing with Item for M 1.
For completeness, we report here the full solution, as a function of the conformal time -= cbit) of Eq. (25). We define
xc 2 1, + 1 2 a+1 :
fv)
6²
(31)
x x+2
then the solution for S(r) given by the inverse function
(-)(w) evaluated at
Sex+ 2Fi
(1,1 + 141:2 x+1 Su(x+1)2
T So
(32)
I + 2
2
as
zos222) 2F
(1, + T+ + 1 S6(x+1)
S(r) f-1
#+
2/T
50
(33)
x x
2
17

---

Grokking at the Edge of Linear Separability
The solution for b(7) is obtained simply by integrating Eq. (25), resulting in
--e0 20l 2F (1.1+
22 st 2t (0(1)))))
1
(1+2)5
b(T)
bog log 1-
I
(34)
x
-000 1-000-5 - 3ol2 Soll-1)azt maz (1,14 12 1te (6(1)))08
(1+x)$-1--0))
242
+ log 1-
(80(2t1)2g [1,1 + 2 + =
x
-0
2+
80(2 x2F
11,1 + 241:2 2+ x+1
a+1 : (00(()))22))
23
80
T 2
2+ x
While these solutions may not necessarily be instructive in this form, appropriate limits can be taken in order to obtain the results in the main text.
E.3. Grokking time in the simplified model
We can define füp feen as the times it would take for the training and generalization loss to reach some threshold e. We can find for by solving Lir Je² (e + (5(1-2))) = We will assume that a is large enough such that S S. from the start (as discussed in the main text, increase the rate that S goes to its final value). Therefore, we plug So log((1 2 and find that for À which is close enough to 1/2, the loss is approximately given by Lir _ Comparing to and plugging the long-time limit b = log( (),, we find that
(35)
ne
Similarly, using the generalization loss Lgon e² we can find that
Igen
2
e (1-2))
(36)
ne
is already clear that for any finite e, egen tú, diverges. We can also obtain an e-independent property by noting that
/logUgen//Gt
log(1 -2X),
(37)
V2
which is verified numerically in Fig. 9.
interesting to note that in the conformal time, we have Rs -T, and therefore can repeat this calculation and obtain
- = - log(ne,, genn 2 log (1 2A) - log 7 2
(38)
In this case, the result is a bit more natural since now the time difference (instead of ratio) becomes e-independent:
gen log2 2 (1 2A).
(39)
We note that this result still depends on E implicitly, in the sense that our assumption that S Soc is true only for long-times or which is small enough.
E.4. Calculation of the subleading term in the separable
We now consider 22 0 but close to zero (that is, are in separable where M -22 is the margin). We know that 1 (1 + M2) + 1] We will now denote
diverges at long times as I+
M
log
M
log (1 + 2²) +1
M
18

---

Grokking at the Edge of Linear Separability
train test
Linear fit
10
23
4.5
vlog(t//t2)
10-39
10 39
10
4.0
10-71
10
3.5
1015
10
1053
10?2
1091
-7.0
-6.5
-6.0
-5.5
-5.0
4.5
t
log(1 2A)
Figure Numerical evidence for the Grokking time in the simplified model. the left panel, demonstrate for the Grokking time is calculated: 4i. ggen are calculated by finding the intersection of the loss with some In the right panel we plot v /log(Lien/ (T)) versus log(1 2A) numerically, and show that result linear with slope J/2 in agreement with the prediction of Eq. (37)
0.001 how
the difference from the diverging term. The equation for is therefore
du " c (-e (u+ M log/ (1 Me M bu(u (a+ 1455 log (3(++(311))
M
(-e
Ōt
2
1 + M t
Plugging the (long-time) solution for the bias, b = 1r241 log [1 M2)t + - get
1
(where by = 0 in our case),
Du
21 T0u (1
M +++
M1M-1 1+M2
(E (e-Na
Nu 1)
M M
Øt
(1
2
For M Rz O, note that the second term is O(M²), and by neglecting it we get
x7
log
2 M ((++M))
M + M +1
For t + oc, and neglecting the other O(M?) terms, we finally get
21
u log
41
82
Remarkbly, this is identical to the result of the in the x2 0 case.
F. Impact of different parameters
Here we present supplemental results for Sections Sec. 4.
F.1. The Variance Scale a
Here we provide additional information regarding the effect of a different than 1. In particular, we will show that increasing a can make grokking more apparent (but only up to certain point). We will first assume that o = 1 at the start, and investigate how taking 2; = aæ; changes the dynamics comparison to that case. We will begin with the non-separable
19

---

Grokking at the Edge of Linear Separability
case (A < 1/2). Recalling that the equations for gradient flow in our model are given by Eq. (7), this results in
as
01 =
0 e Σ St N
bb eeb 0t
=
Nee N
Σ 2 ou;
(40)
We now absorb a into S by denoting 5 oS and investigate how it affects the dynamics of S, and the generalization loss and accuracy as a function of S. First, the GD equations become
o$ Ot
bb Ot
o²
6 N
(41)
We note that the generalization loss and accuracy in Eqs. (3) and (4) are the same except they are now function of |
instead of ||S|] (being a function of a ||]]) Since the equation for as is just multiplied by factor o², the limiting value of So would be the same for the a 1 case, but it will reach it at a faster rate. To sum up, obtaining the dynamics of the loss and accuracy when is larger than one can be done by using the same Eqs. (3) and (4), but also (A) Increasing the starting condition of So by a factor of o, and (B) Multiply only the learning rate of the spatial part by a factor of o², If o is large enough, we can go to the fixed point of ||S] as fast as we want, enabling the appearance of Grokking (if also the limiting value of ||S| is large, which happens when we are on the edge of being separable).
Finally, we will also investigate the effect of a in the separable regime (A > 1/2). Now we can use ?? and Item 2, where we S only need consider how a changes the margin M. Since it is obtained from the equation |||] nm -M, we can
that the new margin will be larger by than the old one, i.e., M o.M Plugging this in the expression for the accuracy in Prop. 3.2, we get that the accuracy is now
1 1 +
lim t1000
Agen
R 2
crf erf
(42)
@2M
where we note that the argument inside the erf is smaller in a factor of o², drastically reducing the limiting accuracy.
F22 Optimizer
The effect of changing the optimizer to Adam is demonstrated in Fig. 11. We note that the fact that adaptive-type optimizers change each learning rate individually based on past gradients, leads the dynamics faster in the direction of || |||. relatively to b. The fact that it makes ||S||| change faster (and not slower) than is probably related to the fact that S is vector in high dimension: Moving each component of such vector will result in change of the norm in rate that proportional to vd, but this may need further investigation. We will also note that using a different optimizer for the non-separable region where, will lead to different solution than the hard margin SVM, as is also discussed by Soudry et al. (Soudry et al., 2018). This means that the results we developed in the main text will not hold, but we can still expect to obtain accuracy smaller than one since ||||| diverges, as indeed can be seen in Fig. 11.
F.3. Initial conditions
As discussed in the main text, changing the initial conditions can change the monotonicity of the generalization loss and accuracy: See Fig. 12, Fig. 13 below.
F4. Different loss
In the main text, we showed that we can use the exponential instead of the CE loss, since it will converge to it at late times. Here we provide numerical evidence that indeed Grokking could be seen, even when taking from the beginning just the Sf teth. See Fig. 14.
exponential loss L N/4 E,
20

---

Grokking the Edge of Linear Separability
0=
a=
a= 10
10a
Train
Loss
1.0
Accuracy
Test
10
-10
10
10
10a
10a
10
10a
10a
10
10a
10
10*
102
10a
10a
102
10a
10a
10
Epochs
Epachs
Epochs
Figure 10. Gradient descent dynamics for three different values of a, all for X = 0.48. The top panels show the loss and accuracy for the and test datasets, while bottom panels present b and the norm of Except a, the parameters are the same Fig. 1. We can see that increasing makes the grokking more apparent at start, but then saturates (that is, increasing will not increase the "grokking time" anymore).
À 0.05
À = 0.48
=
10-2
Loss
Test
10
Accuracy
Train
0.8
0.6
20
|51
10a
10a
Epochs
Epochs
Epochs
Figure Dynamics, using ADAM optimizer with Py Moreo'h's default parameters. The setup is the same Fig. 1, except for the fact that a 1 now instead of 5. Significant Grokking can be seen even though the value of is not large.
21

---

Grokking at the Edge of Linear Separability
||5o|| 0.1
|| =1
||5 =
10a
101
10a
Loss
10
Train
10
1.0
1.0
1.0
Accuracy
0.8
0.8
0.8
Train
10
10
10
115
109
10
102
10a
10
109
10
10a
10a
10
10c
10²
10a
10a
10
Epochs
Epochs
Epochs
Figure 12. Gradient descent dynamics for 0.48 for three different values starting norm, || Except for this, the setup is the same Fig. 1. We can see that the non-monotonicity of the loss can be affected by the starting condition.
bo 0
bp 1
bp =
1
10a
10a
10-1
Loss
10-4
10
10: 2
Train
10
Accuracy
0.8
0.9
0.6
Test
48
10
||il
0-
01
-10
-10
10a
10a
10a
10
Epochs
Epachs
Epochs
Figure 13. Gradient descent dynamics for À 0.48 and for three different values of b. Except for this, the setup is the Fig. 1.
22

---

Grokking at the Edge of Linear Separability
0.05
0.48
=
10a
Train
Loss
10
Accuracy
Test
0.6
10
10
10"
10
10a
10
10
10"
10a
10a
102
10'
10"
10a
10P
10a
102
10a
10a
10a
Epochs
Epochs
Epochs
Figure 14. Gradient descent dynamics for setup which is the same as Fig. 1, but with the exponent loss C = the loss is strictly the exponent loss any time (and not just converge the exponent loss long times, as the CE loss). Clearly, we see that the behavior of the Grokking is similar.
E,
n 4. That is,
N
23

---

Grokking at the Edge of Linear Separability
G. Different input data distributions
As discussed in the main text, our results hold for any data distribution that is symmetric around the origin. Since the underlying mechanism only requires that the data is on the verge of separability (in which case 20/ diverges). As we discuss in App. A, in the limit d, N oo, À 1/2 the critical value below which the dataset almost surely inseparable. Therefore, the analysis and resulting behavior, including the occurrence of Grokking and the critical point of À should hold for any symmetric distribution.
To demonstrate this, we compare three input distributions at À 0.45 in Fig. 15: (1) The isotropic Gaussian input (as discussed in the main text), (2) Non-isotropic Gaussian inputs, generated using a covariance matrix with eigenvalues that follows the scaling law X, 2e with a 1.5. (3) Mixture of Gaussians W(u = +1, a 0.25). We notice that in this context (and its effect on Grokking described in Sec. 4) could also be easily generalized for any distribution, by simply
multiplying all of the inputs by a factor of
Isotropic
Power-law
Mixture of Gaussians
1.0
1.5
1.0
Train Test
1.0
Loss
0.5
0.5
0.5
0.0
0.0
1.0
1.0
1.0
Accuracy
Test
0.8
0.8
0.8
0.6
0.6
10
10
10
0
0
-10 10
-10
-10 10
10
105
10
103
10a
10
10
Epochs
Epochs
Epochs
Figure 15. Grokking for three different input data distributions: Isotropic Gaussian (left), Gaussian with covariance whose eigenvalues follow power-law scaling (middle), and uniform distribution (right). The parameters are: d = 180, N = 400 (A = 0.45), and the optimizer is gradient-decent. Left panel: Isotropic Gaussian with 5. as appears in the main text. Middle panel: Gaussian with cigenvalues that follow An. Ao/n", where
1.5. The normalization factor Aa is chosen such that Σn, An = a d, where here 10. Right panel: each element in the input vector is chosen from mixture of Gaussian distribution, with (4,2 1 and σ1.2 0.25. After sampling, the input was multiplied by 5, as generalization of the original a.
H. Reduction to d-dimensional model with bias.
In this section we will prove the equivalence between the d+ -ddimensional wwo-Gaussians model and the d-dimensional model with bias. Consider the d 1-dimensional classification problem with no bias:
(yitt, 22,2; 04,3, Fi.d+))
(43)
where yı = 1 are the labels, μ is some constant and Ti.j ~ _N(000B). The loss is L $ E, 4 where e, = Ely,S - xx;) is some loss. Taking z; - y; we get
2;; = (H. M...,,,....)
(44)
24

---

Grokking at the Edge of Linear Separability
but the distribution of yo stays the same since it is symmetrical, so we can assume WLOG that y; The dynamics is
ās
aL - Σ
S mz) as mx)) w;.
(45)
We now denote ã; 1 mi 5 ju (so that S a; = S i;), and get
as
0P(S vz) Ti. 0(S xi)
Öt np? Σ
(46)
which is the same except for n² learning rate factor. Since i. = (1, n 11,2 # (..3.....) we can define also the scalar and the d dimensional vectors S, क., by
$ (b, Sa), x; = (1, 2i Fidd)
(47)
We can see that
m 1
0E(S Σ. b)
(48)
Ōt
N
0(S 8. +)
0s
1
Ol T. +b)
mp
(49)
N
Which is exactly the dynamics of the d model with bias, where all of the labels the same. We also note that /B/P where a is factor of the std of T.
I. Extension of the model from another perspective
In this section, we begin with the constant label model and extend it explicitly to a model where only a fraction r of the samples are the same (so that for still be observed. Specifically, we classify point æ; as label -1 if its first coordinate æ;,1 exceeds a threshold μ., and as label 1 otherwise. Here, μ Q(r), where Q is the Gaussian quantile function (inverse CDF) and E [0, 1] is fraction. We note that for 1, only a fraction r of points are labeled -1, while the remaining
we get the same model studied throughout the paper). We show that grokking can
are labeled 1.
The same reasoning as the previous sections can be applied here: near the critical point X = 1/2, the model could reduce the loss first by changing the direction of S, and only then modifying b, leading to delayed generalization. However, unlike the constant-laheling case, full generalization is not possible even for À > 1/2, since the optimal solution not obtained by taking the separating hyperplane to infinity. We can see this by studying the expression for the generalization accuracy (see App. I.1) given by
Agen
=
1+
x
dæie 2 signiiign 11)
2
2
(50)
+ b/S
xerf
V2
VV Ei (5
For perfect generalization, must have b/S1 -1 and S₁ S, Vi 1, which only be achieved up certain error, determined by the max-margin solution.
In Fig. 16, we demonstrate numerically that while the limiting generalization accuracy is smaller than one for < 1, grokking in the sense of delayed generalization and a non-monotonic test loss is still present.
I.1. Derivation of Eq. (50)
Starting from the explicit expression for the accuracy
N. (y:0(S - x. + b) + y:)0(-s m b)),
A
(51)
25

---

Grokking at the Edge of Linear Separability
Train Test
2
r=0.95
Loss
0=99
0
1.0
Accuracy
0.8
0.6
Train
Test
10
[S]
-10
10
104
102
103
104
105
Epochs
Figure 16. Grokking at A 0.48, for three different label-fractions: (constant label), 0.99, and 0.95 (represented by three different opacities). For example, for 0.95, approximately 95% of the input data would he assigned with the label -1 and 0.05 with the label 1. The rest of the parameters the in Fig. 1.
and recalling that 2; (N(0, get that the generalization accuracy is given by
/ d' we - f="s-'= [v()))(( x + + (1 y(=)))(-SS x 6)]
Agen
(52)
del((
Plugging the label to function of the first coordinate
0,
æ <
f(x)=
(53)
1,
else
where μ = Q(r) is the threshold (Q being the Gaussian quantile function as discussed in the main text), we get that
Agen
/ II
drye
25 21 x?
drie- 00 01-S1a1
Σ
S - b)
(54)
V(2mo2)a-1
v
22102
-20
j=2
1=2
+ of dxie 0( +
Σ
Sez0 + b)
(55)
i=2
Setting y Ed Syr,, we that y NOO, ).), where o, = σo Σi= S? =0 -f)). Therefore,
1
5rre dxye 2 202 1 y b) / dxge - + + +bb)
Agen
dye
20
/ /2T03
2
(56)
26

---

Grokking at the Edge of Linear Separability
Taking explicitly the integral over y of the two terms and simplifying the result, will finally get
1
dæie 2020 sign(z1 -w)erf
(
Siz + V2oVISS 2
Agen
1 1+
(57)
2
22TO2
J. Direct calculation of the late time behavior of b, S in the separable regime
Here we present a direct calculation for (t) [t](()] (t), b(t) at late times. First, we will work in the conformal time Jor (t'))tt'. As discussed in the main text, when the data is separable we know that in the late time limit $ goes to
certain direction and ||S||| + oc. Therefore, only the points with maximum S" Fm will contribute to the sum in the large
1 Eestur N
(58)
where S" I' is the maximum value that is obtained, is their degeneracy, We note that S x; are all negative, so the maximum are just the points which are closest zero, these exactly the support vectors. To continue, for -> denote
S~ ~ E
(59)
where E is some constant, and f(T) is some function t. Without loss of generality, and for compatibility with the results of the main text, we will define by the equation arguments: We know that
We can now find s(T) explicitly using the following
S am
i ||S |1 ÔT
ST 0S
S
ST ||]
Di.
(60)
||| ØT
Using the approximation and comparing with
@|| 0r
( we get that
- [(7) = ET)
(61)
N E
Solving for ((T), we get
D 1
f(7) log
11 N F2
7 +
(62)
where C₁ is some constant. Therefore, we get
0b
D 1
(63)
θT
I E
T + Ci
and taking the integral over dr get
E=a
D D
i + +.
B(T) R C2
11 N F2
(64)
Recalling that âr = c², integrate to find:
Bt
1
D +1
1
T(1)
DI E2 E2
"N
E
[C2t + C3
D 1 C1.
(65)
IN E
Using 6 log( 5τ ), get for long times that
E2 log(1).
b(t)
(66)
E
Plugging also ||S|| Ef((T) E'lo((),, we can see that
E log(1).
||()I) R3
(67)
E2
27

---

Grokking at the Edge of Linear Separability
Recalling that E this verifies the result of the main text.
K. Supplemental Details of Experiments
In this section, we will provide information regarding the experiments. The results of the left panels of Fig, 1 (and all of the results Section App. F) can be easily obtained on a personal laptop. As for the heavier experiments which are presented in right-most column of Fig. (and Fig. 7):
(I) Calculation of 5o and S I; properties of the distribution. The setup is: N 24000 1, and =930, 990, 1050, 1086, 11 1110, 1134, 1152, 1158, 1 1164, 1170, 1173, 1176, 1179, 1 1182, 1185, 1188, 1191, averaged over 15000 different random realizations, and using the ADAM optimizer (any optimizer will work in the inseparable regime).
(2) Calculation of the margin. The setup is: N = 2400, a 1, and d 1230, 1260, 1290, 1320, 1350, 1380, 1410, 1440, 1470, 1500, 1530, 1560, averaged over 1000 realizations.
(3) The results of the right-middel and right-bottom panels of Fig. The setup N 400, and d =20, 40, 80, 100, 120, 140, 160, 168, 180, 188, 192, 196, 200, 208, 220, 228, 240, 260, 280, 300, 320, 360, averaged over 100 realizations.
28


Reviews:
[{"review_id": "cNoux1oNWo", "reviewer": "Reviewer_fNB5", "summary": "This work investigates the grokking phenomenon where an increase in test performance is significantly delayed behind achieving perfect training performance. It primarily considered the highly simplified case of a linear model with constant labels, which is analyzed theoretically. In particular, the grokking effect is linked to the transition in $\\lambda := \\frac{d}{N}$ from the non-separable case (i.e. the data cannot be divided by a hyperplane from the origin) to the separable case (where dimensions increase or dataset size decreases such that this hyperplane exists). In the separable case, an overfitting solution will be found; in the non-separable case, a good solution will be found; but at the critical point between the two, a good solution will be found, but the time delay (in terms of gradient descent steps) can be arbitrarily long. These results are intuitively illustrated using an example with just two samples. Finally, in Section 5, some consideration given to the discriminative case where the labels are non-constant.\n\n########## Update after rebuttal ##########\n\nI thank the authors for their response. I have read their rebuttal in addition to the discussion with the other reviewers. I agree that the follow-up work discussed may help to extend the slightly limited scope of these findings. Despite the limitations discussed, I still lean towards acceptance and affirm my initial score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive and useful feedback, and are glad that the reviewer tends towards accepting our paper. \n\nThe reviewer's primary concern is regarding the applicability of our setting to other existing models of Grokking in the literature, which we address below, along with other issues/questions raised.\n\n- *“this claim would be far more convincing...”*: While we believe our paper stands on its own (see the arguments at the end), we agree that it could benefit from a more detailed reconciliation with existing models. The reviewer has provided some good examples and directions. We will try to address some of them here, while an extension of this discussion will be added as an appendix.\n\n- *“..one obvious counter example...”*: Our conclusion does not mean that grokking/delayed generalization will be observed only for one ratio of $d/N$, but rather for a range of values *near* the critical point. However, a ***divergence*** of the “grokking time” appears only at the critical point itself. Note that for $N \\gg d$ grokking will not exist.\n\n- *“Figure 6 of [1] appears...”*: We would expect that the grokking time would be *smaller* for a smaller ratio of $d/N$, but such data does not appear in the figure. Also, Thilak et al. only states that slingshots (not grokking itself) happen - it is not clear if there is a full correspondence. Regardless, the slingshot mechanism itself *relies* on adaptive optimizers (and $\\varepsilon$), so it is not an ideal choice for extending the discussion regarding criticality.\n\n- *“being in the 'goldilocks zone..'”*. This aligns well with our results. Notice that grokking will not be observed even for $\\lambda=1/2$ in our case, unless $\\sigma$ is large enough to attract the dynamics toward the memorizing solution before generalizing (see rightmost panel of Fig. 3 in our paper). For more complicated settings, it is likely that other parameters will need to be tuned in order to see the grokking.\n- *Relation to weight decay:* We thank the reviewer for raising this point: It has recently become more apparent that WD is not a necessary condition for grokking, although it can help observe it in certain scenarios, see for example [4]. Our results support this observation: Adding WD will allow us to see grokking to a certain extent even in the non-separable region. We will add a discussion on the effects of WD in our model in the revised manuscript.\n\n- *“..the aforementioned work of [2] uses ...”* and *“Similarly, [5] find...”*: Thanks for highlighting these models, which are good candidates to investigate the criticality under broader settings. We will consider addressing it in our upcoming work, but we believe this is beyond the scope of the current work.\n\n- *Relation to double descent:* Thanks for this important note. We strongly believe such a relation exists, and that our model is a good candidate for investigating it. We are working in this direction and will address it separately in a future work.\n\n- *Response to the question*: Yes. In fact, these results will hold even for non-symmetric distributions, as this may only change the critical value of $\\lambda$.\n\nTo sum up, while our paper indeed does not offer a concrete recipe for how these conclusions extend to other models in the literature, we still believe that the community would benefit from its publication for the following reasons:\n\n1. It provides a novel and analytically solvable model that exhibits grokking. Being minimalist, it highlights the true relationship between grokking and criticality — something that is somewhat obscured in more complex models.\n\n2. In contrast to many of the current beliefs, we show that neither a transition from lazy to feature learning, nor WD are ***necessary conditions*** for grokking. Instead, the existence of a critical point is the necessary and sufficient requirement. If feature learning is required for this transition, these can coincide.\n\n3. There are strong hints that suggest grokking is related to a criticality [1-3]. Our results support this claim, while offering analytical insights to its mechanism (our approach is somewhat similar to [2], but has much more benefits, see our response to reviewer DzGq for more details).\n\n4. We believe that our results could lay the groundwork for revealing the relation between grokking, criticality, and double descent. In fact, we are currently working on these two directions (the relation to criticality and to double descent), and we intend to address them separately in our upcoming works.\n\n[1] Rubin et al., Grokking as a first order phase transition in two layer networks, 2024.\n\n[2] Levi et al., Grokking in linear estimators – a solvable model that groks without understanding, 2023.\n\n[3] Zhu et al. Investigate the critical data size for language model training through the lens of grokking dynamics, 2024.\n\n[4] Prieto et al. Grokking at the Edge of Numerical Stability, 2025.", "claims_and_evidence": "I generally found the paper clear in its presentation, with the theoretical and experimental aspects convincing in their (specific) claims. \n\nMy main critique of this work is with respect to its broader argument that \"the main takeaway from our setup is that grokking happens\nnear critical points.\" I agree that this claim is essential as, without connecting the highly simplified setting studied to standard grokking settings, the significance of this work would be extremely limited. However, I am not convinced that sufficient evidence (either experimental, theoretical, or through discussion) has been presented to defend this claim. \n\nI would suggest that this claim would be far more convincing if (1) some concrete efforts were made to link the papers explanation to more realistic settings (particularly if some approaches to this extension could be implemented); and (2) if a more detailed reconciliation of the claims of this paper were made to the existing literature on grokking. \n\nRegarding (2), one obvious counter example is that of existing cases of grokking in the deep learning literature in which grokking is reported even as the values of $N$ or $d$ are varied. I appreciate that Appendix C already relates to the original Power et al. work, but there are several other works that would also require reconciliation for the authors' proposed explanation to hold in general. For example, Figure 6 of [1] appears to imply that grokking occurs as dataset size varies. [2] argues that \"being in the 'goldilocks zone for data set size' is necessary but not sufficient to see grokking\" and argues for the relevance of alignment between features and the target function.  Several experiments in [3] relate dataset size to grokking. While I do accept that these works may be consistent with the claims of this work, if this is the case, that reconciliation must be made explicitly (see Section 9 of [2] for an example of this). Relatedly, it would appear that since the theory provided is closely tied to the norm of the coefficients, the authors could relate their findings to the relevance of weight decay/L2 regularization, which has been discussed in several works as an important aspect of grokking e.g. [4].\n\n\n[1] Thilak, Vimal, et al. \"The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon.\" arXiv preprint arXiv:2206.04817 (2022).\n\n[2] Kumar, Tanishq, et al. \"Grokking as the transition from lazy to rich training dynamics.\" The Twelfth International Conference on Learning Representations.\n\n[3] Liu, Ziming, et al. \"Towards understanding grokking: An effective theory of representation learning.\" Advances in Neural Information Processing Systems 35 (2022): 34651-34663.\n\n[4] Liu, Ziming, Eric J. Michaud, and Max Tegmark. \"Omnigrok: Grokking Beyond Algorithmic Data.\" The Eleventh International Conference on Learning Representations.", "methods_and_evaluation_criteria": "Given the highly simplified setting, the methods and general approach taken in this paper are suitable.", "theoretical_claims": "I did not carefully check the proofs of the theoretical claims that were provided in the appendix. However, much of the used theory was adapting relatively standard ideas (i.e. the linear seperability of a dataset, which is understood in the standard logistic regression setting) to the grokking setting with gradient flow and seemed intuitively correct. This, combined with the numerical verifications, indicates that the theoretical claims are reasonable.", "experimental_designs_or_analyses": "Yes, all of the experimental results are numerical verifications that appear appropriate, given the theoretical setting.", "relation_to_broader_scientific_literature": "As previously mentioned, this paper would have been greatly strengthened by efforts to verify its relevance to practical settings. While a complete explanation in all settings is not necessary, some concrete steps could certainly have been taken. For example, the aforementioned work of [2] uses a linear approximation to study the lazy regime and links grokking to a transition from lazy to rich training dynamics. Inutiively, some sort of separation of feature learning and a final layers classification would seem complimentary to this work. Similarly, [5] find a (piecewise) linear approximation of a neural network can capture the grokking effect which may offer a path towards extending the theory to full neural networks. It would seem that these linearized versions of neural networks would make a good starting point for relating this work to more practical settings, even if that link is primarily empirical.\n\nAlso, given that this work links grokking to the interpolation threshold as determined by $\\lambda$ I think it would be valuable to provide even a brief connection with some of the work in the double descent literature. One particularly relevant example would be [6] where the same quantity is studied ($\\gamma$ in that work). Several works have previously attempted to connect the two phenomena (e.g., [7,8]), and it would appear that there are some valuable connections to be made. \n\n\n\n\n[2] Kumar, Tanishq, et al. \"Grokking as the transition from lazy to rich training dynamics.\" The Twelfth International Conference on Learning Representations.\n\n[5] Jeffares, Alan, Alicia Curth, and Mihaela van der Schaar. \"Deep learning through a telescoping lens: A simple model provides empirical insights on grokking, gradient boosting & beyond.\" Advances in Neural Information Processing Systems 37 (2024): 123498-123533.\n\n[6] Hastie, Trevor, et al. \"Surprises in high-dimensional ridgeless least squares interpolation.\" Annals of statistics 50.2 (2022): 949.\n\n[7] Davies, Xander, Lauro Langosco, and David Krueger. \"Unifying grokking and double descent.\" arXiv preprint arXiv:2303.06173 (2023).\n\n[8] Huang, Yufei, et al. \"Unified view of grokking, double descent and emergent abilities: A comprehensive study on algorithm task.\" First Conference on Language Modeling. 2024.", "essential_references_not_discussed": "See the previous section.", "other_strengths_and_weaknesses": "Strengths:\n\n* I found the exposition of this paper to be quite good in general.\n* The figures were well made, clear, and were helpful in illustrating the points they intended to make. \n* The writing was quite clear throughout, including when explaining the theoretical aspects of the work. \n* Efforts were made to provide intuition around the results, which should make the paper somewhat more accessible to a general audience.\n* The paper followed a clear logical flow.", "other_comments_or_suggestions": "Minor:\n\n*L408 should have no parentheses around (Schaeffer et al., 2023).", "questions_for_authors": "The results in this work hold \"for any data distribution that is symmetric around the origin\". Could the authors clarify if this implies that these results, therefore, apply for _any_ data distribution once the input distribution is standardized appropriately as we might do in practice?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Z9Ely8MAKD", "reviewer": "Reviewer_hzJV", "summary": "This paper considers Grokking phenomenon on a simple binary logistic regression model. The authors consider gradient descent (GD) for solving a binary logistic regression problem and analyze the relationship between separability, generalization and overfitting. \n\nIn particular, they consider the case when the input data is generated from Gaussian distribution with labels being the same and the optimizer is gradient descent, what the training dynamics of certain quantities, such as generalization loss and accuracy, will behave under different conditions. They first prove the generalization is equivalent to separability, and then connect the separability with fixed $\\lambda = d/N$ as $N, d\\rightarrow \\infty$ (where $N, d$ are number of data points and dimension) and reveal the grokking happens near $\\lambda = 1/2$.\n\nThen they discuss Grokking in discriminative labeling case (non-constant labels), and finally wrap up the paper by some discussions on limitations and future work.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the reviewer’s thoughtful feedback and positive appraisal of our work. We hope that by addressing the reviewer’s main concerns, they will be able to raise their confidence in accepting our work.\n\n**Weaknesses:**\n\n1. *Limitations of simple models* - The reviewer's main concern is the applicability of our setting to more complex models. Our main claim is that our model contains interesting insights regarding the relationship between grokking and ***criticality***. In fact, the value of our model lies in its simplicity and the fact that it is analytically tractable, allowing us to ***isolate*** the mechanisms that cause grokking which are significantly harder to disentangle in more complex settings.  \n\n\n    Naturally, the next step will be to extend our results to other models. A characterization of the exact relation between criticality and dramatic phenomena, including both grokking and double descent in complex settings, is currently a work in progress, and will appear in our future work. Nevertheless, we believe the community will benefit from the paper in its present form: please see also the arguments at the end of our response to Reviewer fNB5 (starting from: “To sum up, while our paper...”).\n\n2. *Experimental results for non-linear models*: We agree that finding nonlinear models where our criticality results still hold would be interesting. However, we believe that extending our criticality analysis to general, nonlinear models warrants a separate work. The goal of this paper is to show a first, solvable model where the key features of grokking are clearly manifest (non-monotonic test loss and delayed generalization), in isolation from other mechanisms, such as weight decay or feature learning. Nonetheless, there are already many works on non-linear models acknowledging that the ratio of samples to dimensions is a crucial factor for grokking. See, for example, Appendix C, where we discuss the relation to canonical examples.\n\n**Comments and suggestions:**\n\n- Suggestions 1, 2: We thank the reviewer for both of these useful suggestions — we will incorporate these changes into the revised version.\n\n**Questions:**\n\n1. *“Part of the proof is based on the fact that gradient descent dynamics...”*: Thank you for bringing up this interesting point. Empirically, we found that in our model, the maximum value that allows generalization is approximately $\\eta = 1$, and that below this threshold our gradient-flow calculations fit quite well. Nonetheless, it could be interesting to investigate the effects of large learning rates and how they interact with the grokking phenomenon in this model.\n\n2. *“The discussion in Section 5 is good... extend to settings with non-constant labels”*: It is indeed possible to extend our results. As is stated in Section 5, for non-constant labels, grokking in the sense of delayed generalization ***is*** observed, but not in the sense of perfect generalization. However, to demonstrate that grokking to perfect generalization does not require constant labels, we will present a ***balanced-label model that exhibits grokking***. Suppose the first coordinate $x_1$ is drawn from a Gaussian mixture with means $\\pm \\mu$ and variance $\\sigma_1^2$, while the remaining coordinates $x_{i>1}$ are independently drawn from $\\mathcal{N}(0,\\sigma^2)$. Labels are determined by $\\mathrm{sign}(x_1)$. Choosing a sufficiently small ratio $\\sigma/\\sigma_1$ yields arbitrarily large grokking times near the critical point $\\lambda = 1/2$. Numerical results demonstrating grokking in this model can be found here: https://imgur.com/a/4BuMnls. Note that this also proves that the crucial point for grokking is the proximity to the critical point and not the constant labels, which is a good point raised by Reviewer ZiSB.\n\n\n\nFinally, we wish to highlight that our results hold not only for Gaussian inputs but for almost any data distribution (see Appendix H).", "claims_and_evidence": "The claims are mostly for a simple binary classification problem with toy settings, and they are well supported by the proof.", "methods_and_evaluation_criteria": "The proposed methods are mainly some theoretical analyses of a simple model, and the evaluation criteria make sense.\n\nHowever, it is unclear how to extend the analyses to more complex deep learning models, and in the paper there is no empirical investigation on this either.", "theoretical_claims": "I checked the correctness of part of the proofs, which are standard techniques such as gradient flow, inequalities, etc.", "experimental_designs_or_analyses": "The experimental designs are sound.\n\nThe proposed methods are mainly for simple models from a theoretical perspective, and thus empirical verification can be focused on simple settings.", "relation_to_broader_scientific_literature": "It is closely related to the machine learning theory community. The Grokking phenomenon itself is important, and is also related to edge of stability and catapults mechanism when learning rates are large.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Stengths:\n\n1. The authors provide solid theoretical analyses of Grokking phenomenon on a simple linear model (Gaussian input, constant labels, and binary logistic regression) with gradient descent training dynamics.\n\nWeaknesses:\n\n1. This paper only considers simple models and the theory might not be extended to complex neural networks easily. As also mentioned by the authors in Section 6, they did not provide results on nonlinear logistic regression or even more complex deep learning models. The analyses are standard in theory community.\n\n2. The experiments are mostly done under simple settings. It would be better if the authors can provide experimental results on some non-linear models that behave similar as what the theory suggests (e.g., different behavior when changing $\\lambda$).", "other_comments_or_suggestions": "1. Since the proof is mainly based on gradient flow, it might be better if the authors explicitly state it in Propositions? For example Prop 3.2 does not mention this, although they briefly mention it in Proof of 3.2.1 and 3.2.2. Discrete dynamical systems (gradient descent) can be quite different from continuous dynamical systems (gradient flow).\n\n2. It might help if the authors can provide some definitions of certain terms in a separate and formal Def. For example having a separate Def for 'perfect generalization' might be better.", "questions_for_authors": "1. Part of the proof is based on the fact that gradient descent dynamics can be analyzed through gradient flow. I am curious if it is possible to extend the current techniques to large learning rate settings, and thus connect the current framework to other settings like edge of stability and catapults.\n\n2. The discussion in Section 5 is good. I am wondering if it is possible to extend the current proof to settings with non-constant labels. If not what would be the main challenges.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "No concerns"}}, {"review_id": "BqRaQfnQgL", "reviewer": "Reviewer_DzGq", "summary": "The paper develops a minimal setup under the binary logistic classification task to theoretically characterize when and how grokking occurs. They provide both empirical and analytical insights into the mechanism of grokking. The theory utilizes past work on the implicit bias of gradient descent. The paper demonstrates that when the training data is highly unbalanced and on the verge of being linearly separable (from the origin), logistic regression can exhibit grokking.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed and valuable feedback. \n\nThe reviewer's main concern is the simplicity of the setup and how our results could be generalized to other models (a point also raised by some of the other reviewers). However, the reviewer also noted that *“Nonetheless, I believe this is a good first step and should not be used as a basis to reject the paper”*. We believe that we have thoroughly addressed the remaining concerns below (please let us know if you think otherwise). In light of this, we hope the reviewer will consider raising their score. We will now address each point in detail.\n\n- **Relation to broader literature:**\n\n    *Novelty with respect to other analytically tractable models* - We apologize for not clarifying the difference between our work and previous works. Concretely, the papers cited in the *Related works* section which pertain to solvable models **do not** fully analytically solve the dynamics of the models under investigation and do not focus on ***criticality*** as a key aspect of the work. Typically, the results are semi-analytical (using DMFT or related techniques), with the only exception being Levi et al. (2023), which we will address in details below. We will revise this section to include this explanation in the final version.\n\n   In contrast to previous works, our model, being simple enough for analytical investigation, allows us to link the fundamental cause of grokking to the existence of ***critical points***. While we cannot yet prove this rigorously, we conjecture that grokking is intimately related to such critical points in settings beyond simple logistic regression. Most other grokking models are more complicated, making the underlying mechanism more elusive. For example, up until recently it was believed that feature learning and weight decay are necessary conditions for grokking, and our model shows this is not the case.\n\n- **Weaknesses:**\n\n    *The simplicity of the model*:  Rather than attacking the problem  from the direction of the ongoing vast research on the canonical examples of grokking (i.e., modular arithmetic), we try to approach it from a different direction by first analyzing its simpler manifestations. As such, we see the simplicity of the model as an advantage rather than a limitation. Of course, the next step must be trying to relate other examples of grokking to the same insights regarding criticality obtained from the simple model. However, we believe that this is a solid starting point, and that the community would benefit from the paper as it stands. We are currently working on using the criticality approach to study other models and also relating it to double descent — these would be addressed separately in our future works.\n\n- **Questions:**\n\n    1. *Differences from Levi et al, 2023*: We thank the reviewer for highlighting this important point. We are very familiar with the mentioned paper. The main and most fundamental difference between the two works is that, while Levi et al observe delayed generalization near the critical point, their setup does not display grokking in the regular sense of a transition from a memorizing to a generalizing solution. Simply speaking, their test loss does ***not*** exhibit non-monotonicity.\n\n        This limits the scope of their work, since non-monotonicity in the test loss is a hallmark of most known grokking examples. In contrast, our model is the first clean minimal example of grokking that naturally presents grokking both in delayed generalization ***and*** a memorization-generalization transition (see for example Fig. 5 in the appendix, where the non-monotonicity is more apparent since the y-axis is not in log-scaled).\n\n        Another obvious difference is that our model deals with cross-entropy loss rather than MSE, and as such, the dynamics and convergence to a solution are vastly different.\n\n        *It is, however, interesting to note that their result is also closely tied to criticality*: Taking a bit of inspiration from phase transitions in physical systems, we suspect that these two models may belong to ***different*** “universality classes” which are classes of transitions that all behave in the same manner in the proximity of the critical point (e.g., have the same critical exponents). The fact that these universality classes are affected only by “general” properties of the model (for example, its symmetries) makes their identification important, as they imply that one could deduce properties of a complex system by studying a much simpler model, as long as it has the same fundamental properties. While further research is certainly needed to verify this, it is part of our motivation for investigating these simpler models.\n\n    2. *Typos in the example in Sec. 4*: Thank you for pointing these out, both are indeed typos that will be fixed. \n\n- **Comments and suggestions:**\n\n    We appreciate all three of your suggestions and will clarify these points in the next revision.", "claims_and_evidence": "All of the claims made in the paper are well supported by theoretical and empirical evidence.", "methods_and_evaluation_criteria": "Yes the proposed methods and evaluation criteria make sense for this application.", "theoretical_claims": "I did not check the proofs carefully, but all the theoretical statements seem to be correct.", "experimental_designs_or_analyses": "Yes all the experiments are sound.", "relation_to_broader_scientific_literature": "I actually do not understand how this fits into the broader scientific literature. Many of the paper mentioned in the related works section seem to be studying grokking looking at models that are amenable to analysis. What new insight does this model/setup provide?", "essential_references_not_discussed": "I am not an expert in this area but I believe the related works section is quite thorough.", "other_strengths_and_weaknesses": "Overall the paper is quite easy to read and the ideas are explained very clearly. The only weakness is that the setting is extremely simple and it is unclear how we can connect this to even single hidden layer neural networks. Nonetheless, I believe this is a good first step and should not be used as a basis to reject the paper.", "other_comments_or_suggestions": "#### Minor:\n- In Equation 3 it is a bit confusing to use the variable $y$ as a r.v. since the labels $y_i$ are just $-1$ for all $i$.\n- Perhaps it would make more sense to state Prop. 3.1 as $b/\\sigma \\|\\mathbf{S}\\| \\to -\\infty$ to make the connection to Fig. 1 more clear. Or update Fig. 1 to plot $\\|\\mathbf{S}\\|$ instead of $\\sigma\\|\\mathbf{S}\\|$.\n- In the paragraph after the proof of Proposition 3.1 it should be clear that sub-optimal generalization implies \n$c<\\lim_{t \\to \\infty} \\mathcal{A}_{\\text{gen}}(\\mathbf{S}(t),b(t))<1$ it needs to be lower bounded by a constant. Only having an upper bound of 1 does not indicate that it will not generalize. This is made clear in Proposition 3.2 but should be included in this section too.", "questions_for_authors": "1. One major concern I have about the paper is how it differs from other works which shown analytically that grokking can occur in simple models. For example: https://arxiv.org/abs/2310.16441. How do the results in this work differ? Does this work provide a better characterization of how the data influences whether or not grokking will occur?\n3. I'm confused about the example in Section 4. On line 352 and 353. If $\\lambda > 1/2$ then the data IS separable (from the origin) because both $x_1$ and $x_2$ are negative. In contrast, if $\\lambda < 1/2$ then $x_1$ is negative and $x_2$ is positive so the data is NOT linearly separable from the origin. Is there a typo here? Same with line 369/370, isn't $\\lambda > 1/2$ the separable case?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "U0HtcJUrHD", "reviewer": "Reviewer_ZiSB", "summary": "The paper studies the grokking phenomenon and points out that grokking occurs near the critical point where data separability transitions. Specifically, the authors consider a simple logistic regression problem in the limit as the number of data points and the dimension of the model go to infinity. They show that there is a flat region in the loss landscape near the critical point, which causes delayed generalization.\n\n## update after rebuttal\nThe authors' response was satisfactory, and I decided to maintain the score.\"", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the reviewer’s positive evaluation and are pleased that they found our perspective on grokking to offer interesting insights. We have thoroughly addressed their concerns below. If they find our responses satisfactory, we hope they will be able to raise their confidence in accepting our work.\n\n\n**Weaknesses**\n\n- *\"The constant label is restrictive”*: This appears to be the reviewer’s primary concern, so we will address it thoroughly. We believe this restriction is not as limiting as it may initially seem:\n    \n    First, in Section 5, we demonstrate directly that even when the labels are discriminative, grokking can still be observed in the form of delayed generalization - there is a transition from memorization to generalization, albeit imperfect generalization.\n\n    Second, we believe that our model represents a *broader class* of models that can exhibit grokking for similar underlying reasons. Our focus on the specific case where all labels are the same was primarily motivated by its simplicity and the fact that it presents the phase transition in the cleanest manner.\n\n    To further support this claim, we will briefly describe a related sparse feature classification model that also demonstrates grokking behavior near the critical point with ***balanced labels***: Suppose that the first coordinate of the input, $x_1$, is distributed as a mixture of Gaussians with means $\\pm \\mu$ and variance $\\sigma_1^2$, while all other coordinates $x_i$ for $i>1$ are drawn independently from $\\mathcal{N}(0,\\sigma^{2})$, as in our original setup. The label of each data point is determined by the sign of the first coordinate. The model is again logistic regression, but with *no bias term*. By choosing the ratio $\\sigma/\\sigma_1$ to be sufficiently small, one can observe arbitrarily large grokking time provided we are near the same critical point as our original model, at $\\lambda=1/2$. The mechanism behind this behavior is essentially the same as described in the paper: the model initially converges to a memorizing solution (i.e., learning a random separating hyperplane) and only later transitions to the generalizing solution. A numerical demonstration of grokking for this model is included here: https://imgur.com/a/4BuMnls.\n\n    Please let us know if you would like additional analytical or numerical results regarding this model — especially if you believe it could affect your evaluation of the paper. Regardless, we will add to the appendix a brief discussion of this model.\n\n- *“Rigor of the results”*: We believe that our results are rigorous, note that some complementary proofs are left to the appendices. If any specific issue needs further clarification, please let us know and we will gladly address it.\n\n\n**Questions**\n\n- *“Why Adam works for $\\sigma=1$*: We thank the reviewer for this interesting question. First, we will note that numerically one could see the exact same grokking behavior (for $\\sigma=1$) not only for Adam, but also for much simpler optimizers such as SignGD (to which Adam behaves similarly, after some time, due to its ***adaptive*** nature). Therefore, we need to gain intuition for why $|S|$ grows faster than $b$ in SignGD, leading to grokking. \nWe have $\\frac{\\partial|S|}{\\partial t}=\\frac{S}{|S|}\\cdot\\frac{\\partial S}{\\partial t}=|\\frac{\\partial S}{\\partial t}|\\cos(\\alpha)$, where $\\alpha$ is the angle between $\\frac{\\partial S}{\\partial t}$ and $S$. Noting that $\\frac{\\partial S}{\\partial t}$ is a vector of $\\pm \\eta$, we have $|\\frac{\\partial S}{\\partial t}|=\\eta\\sqrt{d}$, where $d$ is the dimension. In the separable case (or on the verge of being separable), after some time the direction of $S$ saturates so we expect that $\\alpha$ would be small. We thus get that $\\frac{\\partial|S|}{\\partial t}\\approx\\eta\\sqrt{d}\\gg\\eta=\\frac{\\partial b}{\\partial t}$. To sum up, this happens due to the adaptive optimizer nature and the high-dimensionality of the model.\n- *Linearly separable data with discriminative labels*: For $\\lambda<1/2$ but using discriminative labels, the model will exhibit late generalization even though the final accuracy differs from 1. For $\\lambda>1/2$, similar to the constant-label case, the dynamics will lead us only toward the memorization solution, making the accuracy level stay at values close to its original value of $\\approx 1/2$.\n- *“Which condition is essential for grokking to occur: linear separability or the labels being almost constant?”*. In our specific model, being close to the critical point with discriminative labels will result in delayed generalization but not necessarily full generalization. However, more generally, ***the proximity to the critical point*** (separability) is the essential point and ***not*** the constant labels: the grokking in the balanced-label model that was presented above is a clear demonstration of this fact.", "claims_and_evidence": "Claims are supported by theoretical results and / or numerical experiments.", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "I did not rigorously check the correctness of the proofs but the results seem to be reasonable.", "experimental_designs_or_analyses": "N/A", "relation_to_broader_scientific_literature": "Grokking is an interesting phenomenon that has not been well understood while a lot of research has been done in recent years.\nThe paper offers a new perspective on grokking, which provides an interesting insight into the community.", "essential_references_not_discussed": "Nothing in particular", "other_strengths_and_weaknesses": "- strength\n    - Grokking is an interesting phenomenon that has not been well understood.\n    - The paper is well-organized and easy to follow.\n    - The paper demonstrates the grokking phenomenon occurs in a simple logistic regression model.\n- weakness\n    - The paper focuses on the case where all inputs are assigned to the same label, which is very restrictive.\n    Although the authors discuss the non-constant case in Section 5 but the condition is still restrictive.\n    - Theoretical results is almost qualitative and qualitative results are provided for the simplified model.", "other_comments_or_suggestions": "Nothing in particular", "questions_for_authors": "- Could you provide intuitive explanation on why grokking occurs even for $\\sigma = 1$ when using Adam?\n- What happens in the case where labels are not constant but the data is linearly separable?\n- Which condition is essential for grokking to occur: linear separability or the labels being almost constant?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors present a detailed analysis of how grokking occurs near the threshold where the number of samples is equal to the input dimension.  They provide a precise analysis of when a solution is merely memorizing vs when it generalizes for a linear regression problem with Gaussian inputs.  The reviewers were largely in support of the work, with multiple pointed suggestions for how to further improve the presentation of the material and its relation to prior work.  I recommend acceptance.
Decision: Accept (poster)