Grokking Beyond the Euclidean Norm of Model Parameters
Authors: Tikeng Notsawo Pascal Junior
Guillaume Dumas
Guillaume Rabusseau
Venue: ICML 2025 poster (2025)
Keywords: Grokking
Delayed Generalization
Regularization
Sparsity
Low-Rank
Overparameterization
Gradient Descent
Implicit Regularization
Primary area: theory->deep_learning
TL;DR: Grokking, sudden generalization following prolonged overfitting, can be triggered using alternative regularizations like $\ell_1$ and nuclear norms or leveraging depth-induced implicit biases without relying solely on $\ell_2$ norms.

Abstract: Grokking refers to a delayed generalization following overfitting when optimizing artificial neural networks with gradient-based methods. In this work, we demonstrate that grokking can be induced by regularization, either explicit or implicit.  More precisely, we show that when there exists a model with a property $P$ (e.g., sparse or low-rank weights) that generalizes on the problem of interest, gradient descent with a small but non-zero regularization of $P$ (e.g., $\ell_1$ or nuclear norm regularization) results in grokking.  This extends previous work showing that small non-zero weight decay induces grokking. Moreover, our analysis shows that over-parameterization by adding depth makes it possible to grok or ungrok without explicitly using regularization, which is impossible in shallow cases.  We further show that the $\ell_2$ norm is not a reliable proxy for generalization when the model is regularized toward a different property $P$, as the $\ell_2$ norm grows in many cases where no weight decay is used, but the model generalizes anyway. We also show that grokking can be amplified solely through data selection, with any other hyperparameter fixed.

Grokking Beyond the Euclidean Norm of Model Parameters
Pascal Jr Tikeng Notsawo! 123
Guillaume Dumas 123 Guillaume Rabusseau 124
Abstract
long training time, it suddenly converges to a solution that generalizes.
Grokking refers to a delayed generalization fol- lowing overfitting when optimizing artificial neu- ral networks with gradient-based methods. In this work, we demonstrate that grokking can be induced by regularization, either explicit or im- plicit. More precisely, we show that when there exists a model with a property P (e.g., sparse or low-rank weights) that generalizes on the prob- lem of interest, gradient descent with a small but non-zero regularization of (e.g., & or nuclear norm regularization) results in grokking. This extends previous work showing that small non- zero weight decay induces grokking. Moreover, our analysis shows that over- pararaaaaaetiiiiiionnn by adding depth makes it possible to grok or ungrok without explicitly using regularization, which is impossible in shallow cases. We further show that the & norm not reliable proxy for gen- eralization when the model is regularized toward different property P. as the &g norm grows in many cases where no weight decay is used, but the model generalizes anyway. We also show that grokking can be amplified solely through data selection, with any other hyperparameter fixed.
Previous work has shown that grokking can be observed by using large-scale initialization and a small (but non-zero) weight decay (Liu et al., 2023a; Lyu et al., 2023). More- over, some works have shown that the lg norm of the model weights can be used during optimization as progression measure for generalization since it generally decreases dur- ing the transition from memorization generalization (Liu et al., 2023a; Thilak et al., 2022; Varma et al., 2023), All these theories have left open the question of whether we always need an lg regularization to observe delayed gener- alization or whether the & norm of the parameter is always good predictor of grokking. This paper attempts answer these questions. We show that the dynamic of grokking goes beyond the Eg norm, that is: If there exists model with property P (e.g., sparse or low-runk weights) that fits the data, then GD with small non-zero regularization of P (e.g., E or nuclear norm regularization) will also result in grokking, provided the number of training samples is large enough and the model is complex enough. Additionally, the regularization of P can be implicit (e.g. model overparame- terization, the choice of training samples). Moreover, the b norm of the parameters is no longer guaranteed to decrease with generalization when it is not the property sought.
We first establish our main theorem (Theorem 2.1), which thcoretically characterizes the relation between grokking and regularization. This theorem is cornerstone of our ar- gument, demonstrating that the generalization delay scales like 1/(a,3), where is the gradient descent step size and 3 is the regularization strength of an arbitrary and appro- priately chosen regularizer that enforces an inductive bias toward generalization. This theoretical characterization ex- tends previous observations, which have been mainly fo- cused on E regularization, providing more general frame- work for understanding grokking dynamics.
1. Introduction
The optimization of machine learning models today relies entirely on gradient descent (GD). The reasons behind the ability of such a procedure to converge towards generalizing solutions are still not fully understood, particularly in over- parameterized regimes. Power et al. (2022) recently ob- served an even more surprising feature of this optimization procedure, grokking: optimization first converges to a solu- tion that memorizes the training data, but after a sufficiently
'nniversité de Montréal, Montréal, Quebec, Canada Mila, Quebec AI Institute, Montréal, Quebec. Canada 'CHU Sainte- Justine Research Center, Montréal, Quebec, Canada *CIFAR Al Chair. Correspondence to: Pascal Jr Tikeng Notsawo <pas- cal.tikeng @ mila.quebee>..
Building upon this theoretical foundation, we validate its im- plications both theoretically and empirically across various settings: sparsity (Theorem 3.1) and low-rankness (Theo- rem 3.4). For sparsity, we focus on a linear teacher-student setup and show that recovery of sparse vectors using gra- dient descent and Lasso exhibits a grokking phenomenon, which is impossible using only Eg regularization, regardless
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Grokking Beyond the Euclidean Norm of Model Parameters
of the initialization scale, as advocated by previous art (Lyu et al., 2023; Liu et al., 2023b). Moreover, with deeper over-parameterized model, there is no need for explicit & regularization, as gradient descent is implicitly biased to- ward such sparse solutions. Similarly, we focus on matrix factorization for the low-rank structure and demonstrate that nuclear norm regularization (denoted (.) is necessary for generalization in the shallow case. This complements prior work demonstrating that deeper linear networks can factor- ize low-rank matrices without explicit regularization (Arora et al., 2018; 2019).
This paper is organized as follows. We begin by present- ing our main theoretical result on grokking time and reg- ularization in Section 2. Sections 3 extend our theoretical framework to concrete sparse recovery and matrix factoriza- tion settings, establishing additional theoretical results and providing empirical validation of the two-phase dynamics predicted by our main theorem. Our findings are extended to general non-linear models and other Nomain-sseecife reg- ularizations in Section 4. Finally, discuss and conclude our work in Section 5.
Throughout this paper, (2-). O(-) and ((-) follow their standard asymptotic definitions.
These findings hold beyond shallow and/or linear networks. We show that l or la can replace lg in a more general setting and induce grokking. We demonstrate this on a non- linear teacher-student setup, on the algorithmic data setup where grokking was first observed (Power et al., 2022), and on image classification tasks. In settings where l regular- ization is not used, the Eg norm of the model parameters tends to grow during training and after generalization, yet optimization still produces a generalizable solution. This directly challenges the previously held belief that the & norm of parameters always a good indicator of grokking.
2. Learning Dynamics: Early and Later Bias
In this section, we will illustrate our hypothesis concern- ing the induction of grokking by regularization. Lyu et al. (2023) proved that enlarging the initialization scale and/or reducing the weight decay delays the transition from mem- orization to generalization, a conclusion already drawn ex- perimentally by Liu et al. (2023a). As we will see below, replacing the weight decay with various other regulariza- tions similarly impacts the generalization delay.
Our contributions can be summarized as follows!:
reg.
mm
mg
(i) We show that grokking can be induced by the inter- play between the sparse/low-ne structure of the solu- tion and the fi/l, regularization used during training, extending previous results on & regularization (Lyu et al., 2023). Our theoretical results extend beyond these specific regularizations, as we characterize the relationship between grokking time, regularization strength, and learning rate in a general setting.
3:22
0.00
3.00
4000
2000
2060
4. reg. d. nom
1. r.gg.. f nom
1. &.rgg.. E. nom
Narr Ox 2o
-2
0.330
0.020
10:
0.350
02
2000
1000
4250
1030
20:1
$.20
Steps 01
Staps
Steps
iii
We show that regularization is necessary to observe grokking on sparse or low-rank solutions. Moreover, we empirically show that in deep linear networks, the sparse/low-rand structure of the data is enough to have generalization without explicit regularization. Adding depth makes it possible to grok or ungrok simply from the implicit regularization of gradient descent.
Figure 1: (Top) Training and test accuracy of a MLP trained on modular addition with l (left), & (middle), and t. (right) regularization for different values of the regulariza- tion strength 3.. Smaller values of 33 delay generalization. (Bottom) In the case of & (top-left), we show the evolution of the E (left), l2 (middle), and f. (right) norm during train- ing. The Eg norm increases despite generalization.
2.001
2=0.01
(iii) Leveraging the notion of coherence, we show that grokking can be amplified through data selection.
Trann
(iv) We show that / can replace le in more general setting and induce grokking. Moreover, in such scenario, and in the shallow sparse/low-rank scenario mentioned above, the grokking phenomenon can not be explained by the &g norm
10036
20003
2204
4000
1003
2000
Steps
Steps
Figure 2: Training and test accuracy of a MLP trained on modular addition with E, regularization for different values of the learning rate a and the L, regularization strength B. When a increases, the generalization delay decreases.
(v) We also show that other forms of domain-spectie specific reg- ularizers strongly impact the grokking delay.
2.1. Motivating Experiment
'Code to reproduce our experiments: https:: htt///gtttuu wwwpuusss/rkkkiig...o.
Consider a binary mathematical operator on S = Z/pZ for some prime integer p. We want to predict y²(x)
norm.
2

---

Grokking Beyond the Euclidean Norm of Model Parameters
T1 o 22 given (x1, v2) 8 The dataset D at x if and only if 4g(x) < x'(yy, x, r). The CL inequality is strengthening of the classical Polyak-Eojasiewicz (PL.) inequality (Chatterjee, 2022), which has been shown to hold for wide overparameterized neural networks in a neighbor- hood of their initialization (Liu et al., 2021). The advantage here is that we will require the CL inequality to be satisfied only at initialization for the theorem to hold, unlike standard results under PL, which require the function to satisfy the PL inequality over its entire domain (Karimi et al., 2020).
((x, x E 5²) is randomly partitioned into two dis- joint and non-empty sets Drain and D the training and the validation dataset respectively. Let us consider as model a multilayer perceptron (MI.P) for which the logits are given by y(x)= = b(2) + w Φ (b(1) + W() (Ee))) o E )), where (ar) stands for the index of the token corresponding to x E S and d(2)= = max(z, 2, 0) is the activation function. E E RRxd1 is the embedding matrix for all the symbols in S. The learnable parameters are E, W(1) E (Raxdd, b(2) E
Theorem 2.1. Take any x(0) Rp with g( ;= -(x(0))
Rdb, (22) E RRxx, and ((2) E RP. We train this model on addition modulo p= = 97 with Irrain == [Panin// = 40%.
O. Assume that 9 is r-CL at x() for some > 0, i.e. 4g
x()) r).. There exist 3max > 0, Olmax > 0 r²x X 1==xxa x(g. and two constants C, C² > 0 such that for all a E (0, amax) and B E (m.)), by defining the subgradient descent update (1) with a and 3 starting at x0), the following hold:
We can see in Figures 1 and 2 that & and & have the same effect on grokking as fg. For all these regularization tech- niques, the smaller a is, the longer the delay between memorization and generalization. Figure 1 (bottom) also shows that in the absence of weight decay, the &g norm of the model parameters is not monotonic after memorization, and even increases without harming generalization performance.
. For any
= 1299)) there exists step
t 2 max (0, 10,-lop log (e/g(0) /log (1 - -O(a x)}
such that g(x("))) S c, [[G(x((1)]1B
O(c) and
x(1) B(x("),)) .r) Vt S t1.
2.2. Theoretical Insights
For any n > 0. mine, (f(x(1) - f*) S 0.5 (n + C'af) 3 if and only i2 > t1 + At(n, (1).
We now provide high-level theoretical insights on how the dynamics of regularized gradient descent can induce grokking. Let g : R - [0, oo) be a differentiable function, h: RP - [0, 20) be a subdifferentiable function and /3 >> O. Typically, g is the loss function of an overparameterized neural network on the training data, while h serves as the regularizer. Our goal is to minimize the composite objective f 1:00 + .hh using subgradient descent with learning rate > 0. The update rule for this problem is given by
with At(n, (1) == (x ,0) Moreover assum- ing , n 0g + D. mine, Sist, (h(x(t)) - 11))) = 0.5 (1 + +Can+ a3) if and only ift > 11 + At t1).
Proof Sketch. We show that, as long as 3 || (x(())))]l S
|| []
YO S k < t, for a certain t 2 1, the follow-
ing hold: x(k) E B(x("), r), (x)( S (1 - (1-5)-g(e 9x
and [
S 2L (k) for all 0 S k S t. for a
certain ō E (0, and L. > O. Using this result, 2(88), constant C > 0 such that for any precision we can adjust hyperparameters to ensure that y reaches this precision before (1]H()) > |[G(-)]l. We then bound f(x())) f" and h(x()) - hn by quantities that depend on 4. a, B. Using these bounds, we extract the delay A1(1, 41) that ensures the desired precision 17 on f(x(t)) - f and h(x(9) - hū after any step f1. A more formal version of this Theorem and the corresponding proof can be found in Section C. of the Appendix.
find
x((+1) = =x(0- - a (G(x(0)) + (H(( Vt 0 (1)
where G(x) = Vg(() is the gradient of at x and (x) E 0h(x) is any subgradient of h at x. We let
f := infxeR" f(x) and θ := arg min fx) C R'. Similarly, we define 9" and θg. and assume g" = 0 with- out loss of generality. Lastly, we define dist(x, θ,) := infyee, ||x yll2. Intuitively, the training dynamics can be decomposed into two phases under certain conditions on a and B. The iterates x(") initially move toward solu- tion close to the initialization (0) that minimizes g (for in- stance, the kernel solution associated with g). Later in train- ing, the influence of (x) dominates the update, driving f (x(()) z fr and ())(t)) h infxee, h(x) within additional @(1/a training steps.
Theorem 2.1 highlights a dichotomy between memorization and regularization and formalizes a two-phase dynamic in the training process, under the CL assumption on the loss function 9 at initialization. The first phase, described in the first bullet point of the theorem, shows that for a suf- ficiently small regularization strength ,, the iterates (!) remain close to initialization and quickly minimize g up to any precision = Q(BC). This corresponds to memoriza- tion of the training data, as the model fits the loss while staying in a local region where regularization does not dom- inate. By choosing 3 too large, we may never minimize g. hence the lower bound S on e, which vanishes as 3 - 0.
We define the Chatterjee-Lojasieviii (CL) constant for at Xx E R" with radius r > 0 as (Chatterjee, 2022):
inf x(e.x,1) r) y )0 || /g(y)1///(>)
(2)
where B(x, r) == (y E Rp| ||x- ylla= S r). The function g is said to satisfy the r-CL inequality at x or to be r-CL.
3

---

Grokking Beyond the Euclidean Norm of Model Parameters
In the second phase, captured in the second bullet point, when g and its gradient are already small, the regularization term 3h(xx gradually drives the iterates toward low values of f and h, until reaching (x((1) = f and h(x(*)) = h up to an error of order (aa:3²) and ((a33) respectively, within At = (1)333) additional steps. If this late-phase bias induced by h is aligned with generalization, then the grokking time is proportional to 1/a,3. This conclusion is consistent with previous findings for h(xx)= [|x] where grokking time scales as 1/3 3 (Lyu et al., 2023). The learning rate a does not appear in that result because Lyu et al. (2023) analyze the continuous-time gradient flow setting.
matrix and & E IRN a noise vector, we want to minimize f(a) (a) + (3((a) using gradient descent with learning
rate a > 0, where g(a) = 2 |a --112 and h(a) =
||||1. Like in Theorem 2.1, the update ( first moves near
the least square solution á (xTx) x y leading to memorization, with g(ā) < 1 | Later in training, @h(a) dominates the update, leading to ||a(1|| - I|a"" || O in the order of 1 1/(a,8) more training steps. When a33 is small, this causes non-trivial delay between memorization and generalization (i.e., grokking), for a suitable X.
Theorem 3.1. Assume the learning rate, regularization coefficient and the noise term satisfy 0 < a (X X) < 2. 0 < 8/ maxx (X and |X Ell2 S |
We also establish that under some assumptions on h and 9, if g is r-CI. x E Rp, then for all E (0, 1), there exists Imax = mma x)
VCaB, C >
0. Let P2 == -max (In - x There exist t1 < 00
0 such that for all B < Smax. the function = 9 + Bh is also er-CL at x, i.e., 4f(x)= < <E²r² x, er). As consequence, the first point of Theorem 2.1 applies to f - f when 3 S 9max More specifically, we have x*)) E B(x("), for all i, and as t + oc, x(1) converges to a point x" E Bxx((), Er) where f(x") = f Moreover, for each t 2 0, [|x(0) - x²1 < (1 (3=6)"=-,² and ((x())-f+ S (1 0)
and C' > 0 such that ||a()) - áll2 S 20 1/2
Vt 2 t1, and
1-pg
ming, (1a(0)) -|aa* 11) S
g -+(C+0")a8
l2 2
4 (1) for any 1) > At(n, (1) = (a -
az9
Proof. See the proof in Section C.4 of the Appendix.
Now let illustrate why Xa = y* (memorization) and ||a| ||| = ||a'|1 are enough to conclude a!*) = a (gener- alization) when N is large enough with a specific class of design matrices X commonly used in practice (Fou- cart & Rauhut, 2013). For that, let us consider the prob- lem in a more general context, where we have measure- ment matrix M E R xn and a list of noisy measurements y = =b (M) + & E R of an unknown signal b' e R" assumed sparse in a know basis &b E R"*", i.c. b' = Σx= a 2., = Qa with ||a"|lo < n. The aim is to find a" by minimizing ||allo subject to ||F (M) - y''ll S C, where an upper bound on the size of the error term &. Elll S Since this is NP-hard (Natarajan, 1995; Donoho, 2006a), ||allo is often relaxed to its tightest convex re- laxation. ||a]| (Foucart & Rauhut, 2013; Chandrasekaran et al., 2012), leading to the convex problem of minimizing ||al| subject to | - y S E, with X = MΦ. This is equivalent to the problem of minimizing f(a) presented above, for suitable choice of B (Rockafellar, 1970; Boyd & Vandenberghe, 2004; Candes et al., 2006).
with 8 = min(1, θ (x(f,x()) a)) The proof of this fact can be found in Section C.3 of the Appendix. The result required that Vx E R", 0 E ah(x) = h(x) = O. A funetion h- R R satisfies this if, for example, it is convex and there exists x* E RP such that h(x') = infx h(x)= > -oo. This is the case of f1/2/> norm.
3. Beyond Weight Decay
In this section, we will illustrate examples in which late phase bias associated with generalization is not driven by fz. namely, sparse recovery and low rank matrix factorization. Indeed, we will see that the use of weight decay alone causes an abrupt transition in the generalization error, as predicted by previous works, but that this transition does not correspond to generalization. We will also show that for a fixed training data size, the depth of the model used or an appropriate choice of training samples can significantly reduce the grokking delay.
We consider the operator Fa (M) = Ma E RN that take N measurement vectors (M, E R"); and return the mea- sures (M. a}x of a E RR" For a matrix A E Rm the operator vec E IRTnn stacks the column of A in a vector; mmax/ min/ // (A) is the maximum/minnmumm/"" singular val-
While classical sparse recovery assumes that the true signal is exactly sparse and the measurements are noise-free free, in practice, it is common to design the measurement matrix M for stable and robust recovery. Stability refers to the ability of a reconstruction method to recover an approximation a of signal a" whose energy is concentrated on a few compo- nents (i.e., approximately sparse), with reconstruction error controlled by the sparsity defect inf Infjalllse ||a* - -alI. the distance from a' to the set of exactly s-sparse vectors. Ro- bustness, on the other hand, means that the reconstruction is resilient to measurement noise, i.c., small perturbations in the observed measurements y = Fb. (M- + g lead to
ues of A; A ||+ = Σ, (A) and ||Al|2-2 == m (A).
3.1. Sparse Recovery
Consider a sparse vector a* E RD, i.e. 8 == ||a" loo < n, where ||a"|lo is the number of non-zero components of a". Given y" XXa + & with X E R.Nx design
4

---

is com- Grokking Beyond the Euclidean Norm of Model Parameters
10
small deviations in the reconstructed signal a, with the error sealing with the noise level
100
1-
10-1
x|a y'1la / lly ll
Definition 3.2 (Robust Null Space Property). A matrix A E [R*x" is said to satisty the robust null space property with constant p E (0, 1) and т > 0 relative to a set S c [n] if || uulli < pl uml \ 1 + - A|/// for all E RR; where us = julles E Rlsi. It is said to satisfy the robust null space property of order s E IN if it satisfies the robust null space property relative to any set S C [ra] with [S] S.
Error
I1a
00 ||a|||
0.6
Gradients
I18H(a"") / IGG("
0.2
100
10a
000
10
100
109
Theorem 3.3. Assume the mairix X E RNNx satisfies the robust null space property with constant p E (0.1) and t > relative to the support of a². Then, under the same condition as in Theorem 3.1 on a, 33 and &. there exist C1 C22 C3 0 such that min,, StSts ||a(0) la!'-a""l alI1 Cyn+ CzaB + CsllEl12 if and only if 22 + At(n, (1), Vn > O.
10-4
10-1
Error
10 2
10-
Ar
10 3
10
||a 1|:
||a
10-1
10a
10c
1302
10a
10
102
100
Proof. Proof in Section C.5 of the Appendix.
Components of att)
0.10
0.08
It is worth noting that we should instead quantify the train- ing error 1 Eim (v(),) - y" (xx)))² g(a) and the gen- eralization error E(a)= = Fx,c (u(x) - vā (x))².
0.06
0.04
0.02
mon in the context of grokking, with v(x) = x and (x) = x3 a + in our case. Assuming Es- = 0, and using == E[xx J. we get E(a) (a- a²) (a a²) + Esa. Random matrices with independent entries, notably Gaus- sian and Bernoulli matrices, allow for the lowest restricted isometry constant and, thus, better recovery of sparse vec- tors (Rauhut, 2010). Under this independent and identically distributed (iid) assumption, we have Σ = o²I,, for a cer- tain a > which implies E(a)= = 0²I|a a 112 + Ee?. So ||a - 112 captures well the notion of test (or validation) error commonly used in the context of grokking, while E captures the irreducible part of that error.
0.00
0.02
10a
10
102
10
10a
10a
Steps (t)
Figure 3: Relative errors, gradient ratio, norm of ||a ln, and components of a!)) as a function of t. G(a(t) dom- inates (a( until memorization at th, g(a(ti))
0.
From memorization (a( dominates and make (|a(1)] converge to ||a" aa'l1 at t2 = t + At, and so a = a".
Validation Experiments Using (,,8,N,a. 3) (10?, 5, 30, 10 10-5), we observe a grokking-like pat- tern, where the training error |Xa(() y" first decreases to 10 then after long training time, the recovery error [t( - l|2 decreases and matches the training error (Figure 3). The generalization results from the interplay between the sparsity level s = ||a" llo. the number of measures N, the & regularization 3 and the learning rate a. Small s requires small N for generalization. Generalization occurs mainly for small (but non-zero) (Figure 4), Large 8 pushes the recovery error to plateau at suboptimal value early in training (causing complete memorization) or oscillations (causing no memorization). However, small values require longer training time to plateau and generally do so at a lower value of recovery error (grokking). We provided more experiments in Section H. 1 of the Appendix.
Memorization When the initialization (1) is close to á, it takes less time to memorize since t1 decreases with a3 [ Another alternative for reducing the term I-p= and guaranteeing perfect memorization earlier is reduce 3. But this increases the generalization delay At. Low signal-to-noie noise ratio |a*) / 2 and reuulaagultiization (large () also make perfect memorization difficult.
Generalization After memorization, when ||a("|]
|| becomes too small, we have ||a(() - a1 (Fig- ure 3) since for the problem of interest, the sparse solution a" the minimum & solution to || ll2 S e under the sparsity constraint = ||a*|lo < n and certain assumptions on X (as illustrated above). The additional number of steps it takes to reach this minimum -n-mo solution is inversely proportional to a,, so that the smaller a is, the longer it takes to recover a" (grokking), and the smaller is the error (ta a" a²1|1 when - oo (Figure 4).
Other Iterative Method for & Minimization The above results hold for the projected subgradient method, for which the update writes (f+1) = IT (a(t) (a))) with
5

---

Grokking Beyond the Euclidean Norm of Model Parameters
necker product) has orthonormal column since & T VV+ & (U*T**) h.- This framework encom- passes several matrix factorization problems. Matrix sensing seeks the matrix A from N measurement matrices (X; E R
10
10a
10-4
10a
10a
10
]E(NN and measures y* = (tr(X] (e)) In [vec(X,)]kEM E RN*= since yi Far (vec(X,)). For a matrix completion task, we have N measurement vectors?
10
10-5
Error
this case X =
10a
10-6
(x(), x(²) E R" x xR" and mea-
4
10a
|----|1-
10
sures y = = A A'X (2) = Fa-
(x( & x that (2)
10
à 1|1
y" = F (X) with X = x x E RN (face-
aB
10-3
splitting product).
10
8
10
7
10
10
10
10 3
10
aB
We now analyze grokking when minimizing f(A) | vec(A) - y"11 + A|| via subgradient descent with learning rate a. Like in sparse recovery, under some condition on a and 3, the training dynamics can be de- composed in two phases; the memorization phase where update A (1) first moves near the least square solution veo(A) A) == (XTx)'XTy² y", and generalization phase where A (9) converge to the minimum E, solution.
Figure 4: Generalization step l2 (smaller & such that ||a(t) a ||22/|a"* I|2 S 10-1) and recovery error [|a((t) a a1I| as
function of aB8. We can see that t2 x | lå-a0 a=/loo/ad /a3 and ||a(t>) - a" ll x a:8, i.e. small a.3 require longer time to converge, but do so at a lower generalization error.
II the projection on the set {a Xa y"); and for the proximal gradient descent method; for which the up- date writes a (t+1) Sae (a - (G)) with (a) = sign(a) @ max([al - 7, 0) the soft-thresholding operator. Note the dichotomy between these two methods. Fach uses one of the terms of fa) as the objective to be minimized and the other as a constraint. After every gradient descent update with the gradient of the objective, the iterate is pro- jected onto the feasible set of the constraint. Still, they all exhibit a grokking behavior (more details in Section D.1 of the Appendix). One training step is enough for the projected subgradient to get zero training error. This further shows that generalization here is primarily driven by h.
Theorem 3.4. Assume the learning rale, the regularization coefficient and the noise satisfy < max X X. < 2.
< B /min(nn; 122) < +max (XX) and |X Ell2 S
VCaß, C <>0 O. Let P2 = -max (%h- 10x X). There
exist li 00 and C" such that || vee(A(t) (1)
A)ll S
2/2
I-p2 Vt 2 t1, and mine (IIA())) ||- ||A S
(C+C")) - 12 2 4 41 for any n > 0, with
2
At( = | 1
agj
Proof. See the proof in Section C.6 of the Appendix.
For a suitable X, X vec(A y (memorization) and A (0) ||ll are enough to conclude A (!) A
3.2. Matrix Factorization
Theorem 3.5. Assume the linear measurement map F. F.(X) satisfies the robust rank null space property of order with constants p (0,1) and >> O, ie for all A E ROLXMS, ((()) S P Emin(mi man ae(A) + T ())) Then, under the same condition as in Theorem 3.4 on a. 3 and &, there exist C2, Og > 0 such that mint, StSta (t) A'll, S Cin+ C2a + CallEl1a if and only if 2 t At(7, (1),, for any n > 0.
In this section, we study grokking for matrix factorization and extend the results of the previous section on sparse re- covery. Given a low rank matrix A" E Ra x of rank r and measurement matrix X E R.Nn. X.x m2; we aim to min- imize rank(A) s.t. [Fa (X) - y2 S c for A E RMixRa x and a = vec(A) E nR; where y* = =F -vec(A-) (X) + E are the measures and & = R.V the error term with ||El|a S e. This is NP-hard (Vandenberghe & Boyd, 1996; Fazel et al., 2004). The usual convex approach for matrix factorization is to minimize ||A|| |l. s.t. ||5, (X) - y S e since the trace norm the tightest convex relaxation of the rank (Fazel et al., 2001; Recht et al., 2010; Candes & Recht, 2012).
Proof. Section C.7 of the Appendix.
Despite the similarity between this result and the one ob- tained on sparse recovery, this similarity is trivial only in the carly phases of training. In fact, in sparse re- covery problems, we take into account only the iterate
Let n = n1n2 and A = V be the full SVD of A We are dealing with a compressed sensing prob- lem with the signal vector a" = voc(A") = ob vec(E*); where vec(E*) E R" is sparse since || | llu = r S min(n1,n2 n2) < n; and + = V s U E Rax*n (Kro-
FFrr standard matrix completion, each x()) (resp. x(2)) rep- resents row & [m]] (resp. a column c E [na] of A' (one-hot encoded in dimensions ni and na respectively), giving y = A
6

---

Grokking Beyond the Euclidean Norm of Model Parameters
a0), whereas, in matrix factorization, we take into count not only the singular value matrix Σ of the iterate A (4) but also its matrix of singular vectors U(O) (left) and ((t) (right). We used the Wedin's 's sin θ bound (Wedin, 1972) to quantify the variations of singular vectors after the memorization phase, and the result shows that they change by at most O (a.3(77 1)) at each iteration, with 7 = min (A))) max (A*). Also, if we take a matrix factorization problem and optimize it with only &,, there is no grokking unless the matrix is extremely sparse so that the notion of sparsity prevails over the notion of rank, which shows the point of studying l. separately. Finally, neural networks trained with l1, lg, and l. have very different properties.
100
2.0
10-1
Gradients Ratio
1.5
10
At
Error
10-3
10"4
||XvecIA (1) -y 112 / 11x
112
||A A'II= 1AA 'I+=
10-5
[|8H(A#)]]== []
10-6
10a
10a
10a
107
109
15
10-1
10
At
Error
10
||
10-3
Generalization When ((A) become negligible compare to (A), the singular values start involving approximately as (A (++1) - -A') )-- = lo, (A - A* - a.3 up to and error of order 0,3/(7- -11] 1). This leads to a generalization through a multiscale singular value decay phenomenon (Figure 5). The small singular value after memorization converges to 0, followed by the next smaller one until ||A (+)) ||A]|] This process requires θ (1/a33 steps.
10
|| |1-
10-5
||A
10-6
10a
10
10
107
Dimensioo min(m,,
Singular values
Validation Experiments We set (n1, 722, T, N, a, = (10, 10, 2, 70, 1001 1011,100 1)_ and optimize the noiseless ma- trix completion problem using subgradient descent. We observe grokking-like pattern, where the training error |X vec A (0) y first decreases 10-4, then after a long training time, the recovery error (I) A de- creases and matches the training error (Figures 5 and 6).
100
22
10
10a
10
107
Steps (t)
Figure 5: Relative errors, gradient ratio, the norm
||
Other Iterative Method for & Minimization The above results hold for the projected subgradient method, for which the update writes A (++1) II (A) (t) aA (*)) with II the projection
and evolution of singular values. G(A (9) dominates (H(A()) (0) until memorization (1 S 41). From memoriza-
tion (H (9) dominates and make | (9I. converge to
the set X vcc A y"); and for the proximal gradient descent method; for which the update writes A (++1)
|| at t t At, and so A = A
Saz (A aGG( ( ())) with S,, = U max (2' - V) the soft-thresholding opera- tor for A = UEVT under SVD, where max(E - y, 0), 5. max(E.; - 7, 0). More details Section D.2.
insufficient training samples (such as sparse recovery or matrix completion using number of samples far below the theoretical limit required for optimal recovery). We call this phenomenon "grokking without understanding" like Levi et al. (2024) who illustrated it in the case of linear classification, and we attribute it to the fact that the assump- tions underlying prior theoretical predictions, particularly Assumption 3.2 of Lyu et al. (2023), are violated in our setting (more details in the Section of the Appendix).
3.3. Grokking Without Understanding
Contrary to prior studies (Lyu et al., 2023; Liu et al., 2023a), we find that in the over-parameterized regime (N < n),, large-scale initialization and non-zero weight decay do not necessarily lead to grokking (Figure 7). In fact, by using only the & regularization under large-scale initialization, there is an abrupt transition in the generalization error dur- ing training, driven by changes in the norm of the model parameters. This transition, however, does not result in convergence to an optimal solution and can arise even in cases where the problem has no optimal solution due to
For the problems above (sparse recovery and matrix fac- torization), by replacing h(() in f(a) g(a) + (3h(a) by h(a) = ||al the model converge to the least square solu- tion à = (x x 3lm) x yTy,, but this solution can not give rise to generalization when N < n.
Theorem 3.6. Define P2 = ||, a (XTX + (3,,, 112
7

---

Grokking Beyond the Euclidean Norm of Model Parameters
e 14 = Tu | T (1 and Vj na VIV (i.j) E [mi] x [m2] quantify how strongly individual rows and columns align with the top singular vectors, where ((m)) be the jth vector of the canonical basis of These quanti- ties, also known as leverage scores, indicate the "influence" of each row i or column j on the low-rank structure. A row/- column with a high leverage score projects strongly onto the span of the singular vectors, meaning that relatively small number of its entries capture much of the matrix's 's structure. Uniformly low coherence (14 and close to 1) implies that the matrix's ' information is well-distributed across rows and columns, thereby reducing the number of samples needed for exact recovery. The significance of local coherence extends sampling strategies and recovery bounds. For ma- trix completion, given N n2 and [0, 1], we select the first NN entries (ž, with the highest values of μ4; + Mj- and the remaining (1 - T)N entries uniformly among the rest. As T + 1, performance improves, and the number of examples required to generalize decreases exponentially, as does the time it takes the models to do so (Figure 8).
|| (for ee (ma)
10
10
10-1
10
10-2
10a
10
10a
Error
10
10a
42
10-5
10a
IA Al
10T5
HIA
-A I1+
10a
10
10
7
10
10
10 4
10 3
10
2
10
aB
Figure 6: Generalization step l2 (smaller such that |A(())
A ||/l|| |r/l||A|| |/ S 10-4) and recovery error || (t2) as function of a,3. We can see that 12 x ||A A 2 and ||
A|
I x aB, i.e. small a3
require longer time to converge, but do at lower gen- eralization error. The outlier for very small a:3 is due to insufficient training.
Unlike matrix factorization, large values of coherence are detrimental to generalization for compressed sensing. The incoherence between the mea- surement vectors (rows of M) and the sparse basis (columns of b)) is crucial for successfully recovering b* Fh from the measures y* = Mb* + &. If M is incoherent with Q, each measurement captures dis- tinct "view" of b", reducing redundancy. This diversity of information allows for the successful reconstruction of a* even with fewer measurements (e.g., below the Nyquist rate for signals). We also observed that higher coherence in- creases the number of samples required for recovery and de- lays generalization, while lower coherence results in faster generalization and better recovery. While grokking has been studied extensively, the impact of data selection on grokking remains largely unexplored, making this one of the first works to address this critical aspect. We provide more details in Section G of the Appendix.
20th
Figure 8: Training and recovery error as a func- tion of data size and co- herence
10
107
Steps 10
Steps
Figure 7: Relative errors and the the norm ||a ll2 for large initialization scale ||a("||2 10 and small weights decay 3 = 10-5, Weight decay alone just reduces the components of the iterate (*)), not la""|o--
Assume the learning rate satisfies 0 < < omar (X
2
X)+3
Then (la al/2 I p5 ()) - -al/2 Vt On the other hand, for N < n, l 1 Z
>
1.
ll (L,, (xXX nonzero component orthogonal to the column space of X, then å cannot perfectly generalize to a".
a" In particular, if has a
Proof. Section C.8 of the Appendix.
Even when l is present, the weight decay 3 is choose such that || < aB, then (1) will get stuck near å, and there will be generalization. So, a bad choice of 3 can be detrimental to generalization (it is better not to use f on that problem unless the initialization scale is nontrivial). Since the minimum l norm solution only gives memorization, the fg norm cannot be used as an indicator of grokking.
3.5. Implicit Bias of the Depth
We explore the role of overpa- rameterization in sparse recovery using linear network parame- terized as i where the depth I, > 2 introduces over- parameterization without altering the linearity of the function class. Our findings reveal that depth can replace E erreguuaaiizttion for generalization when initializa-
8:00
03
06
370
3.4. Amplifying Grokking through Data Selection
Figure 9: Recovery error as a function of depth and data size.
The analysis of recovery guarantees for matrix factoriza- tion hinges on local coherence of the target matrix A' = V*T (compact SVD). The local coherence measures
8

---

Grokking Beyond the Euclidean Norm of Model Parameters
tion is small, as the gradient updates introduce precondi- tioning effect that promotes sparsity and enables signal re- covery. Unlike the shallow case (L = 1), where sparsity can- not be enforced without f 1-regularization, depth provides an implicit mechanism that biases updates toward sparsity, resulting in better generalization with fewer measurements (Figure 9). Additionally, depth reduces the generalization delay, with abrupt phase transitions and a staircase- loss curve during training. The generalization error decreases with increasing depth L, showing that depth effectively com- pensates for fewer measurements and facilitates recovery, albeit at the cost of longer training times. For L 2 2, large-scale initialization combined with small non-zero lz- regularization results in grokking, unlike the shallow case where the phenomenon of "grokking without understanding" observed. More details in Section F.l of the Appendix.
consistent with physical laws. Sobolev training (Czarnecki et al., 2017) generalizes this idea by incorporating not only input-output pairs but also derivatives of the target func- tion. We optimizer the student from the previous paragraph by adding on the objective function the first order Sobolev
penalty 8 N Oye (xx) ay" -==
(xx) ||: where the hyper- 8x
parameter 33 ensures that the model not only fits the data but also respects known smoothness constraints or differen- tial structure, which is crucial in physics based applications. Large a3 values lead to faster grokking (Figures 11 and 34).
a- 0.0001
- 0.001
a- 0.00055
Loss
10
(t)
(1)
We discuss deep matrix factorization in Section F.2 of the Appendix, as it is already well-known that overparametriza- tion by adding depth makes it possible to recover low-rank matrices without any explicit regularization (Gunasekar et al., 2017; Arora et al., 2019; Gidel et al., 2019; Gissin et al., 2019; Razin& & Cohen, 2020; Li et al., 2020).
Figure 11: Training and test error two layers ReLU teacher- student with Sobolev training. for different values of the learning rate a and the Sobolev coefficient, ,3.
Classification On the algorithmic dataset (Power et al., 2022). L and l, have the same effect on grokking as lz, i.c., smaller regularization coefficient (and learning rate) delay generalization (see Figure 2 for la, and 30 for l and fz). We observe a similar phenomenon on a two-layer ReLU MLP trained on MNIST (Section H.3.4).
4. General Setting
In this section, we show that é1, Eu, and other domain- specific regularizers can replace f2 in a more general setting and induce or control grokking.
5. Discussion and Conclusion
This work extends the understanding of grokking, show- ing that the transition from memorization to generalization can be induced not just by f2 regularization but also by sparsity or low-rank structure regularization (e.g.. l and nuclear norm regularization) or domain-specift regulariza- tion. These findings are particularly relevant in practice, where large-scale initialization is not always feasible, yet grokking still occurs. Sparse and low-rank-bard with good generalization performance are very useful in machine learning today, not only because they consume less memory during training and inference but also because they are more interpretable than their dense and full-rank coun- terparts. The sparsity and low-rank assumptions are also central to techniques such as Low-Rank Adaptation of larger deep learning models (Hu et al., 2021) and model pruning (Han et al., 2015) in resource-efficient deep learning; sparse autoencoders in mechanistic interpretability (Cunningham et al., 2023; Bricken et al., 2023) for AI safety (Bereska Gavves, 2024); name a few.
Non linear Teacher-Student We consider a teacher (x) (Ax(A'x) x) from Rd to R with r hidden neu- rons; where x(2) = max(2 (2, 0). We independently sample and op- timize the parameters 0- = (A, of a student (x = Bo (Ax on them with the square loss function g(0) = EN ii= [xx) (x))) and different regularizer ((0). Ep for p E 21. For all of these regularizer, the smaller is a13, the longer is the delay between memorization and generalization (see Figures 10 for the training curve with L, and 31, 32, 33 for more results with f,/2).
N inputs output pair D =
models
0. 0000001
0 0055
Train
Test:
12
Steps 10
Stops C3
Steps 10
Figure 10: Training and test error two layers ReLU teacher- student with & regularization, for different values of the learning rate and the & coefficient B.
Our results highlight that in deep models, gradient descent implicitly drives the model towards solutions with sparse or low-rank properties, effectively mitigating overfitting (Arora et al., 2018). We also study the impact of data selection on grokking and demonstrate that it can be mitigated solely through data selection.
Domain Specific Regularization Physics-Informee Neu- ral Networks (Raissi et al., 2019) leverage prior knowledge from differential equations by incorporating their residu- als into the loss function, ensuring that solutions remain

---

Grokking Beyond the Euclidean Norm of Model Parameters
Acknowledgements
neural network. CoRR, abs/ 011188006005551 2018. URL http://arxiv. org. www./agg////..o.oooooo.........r als
806 05451.
We thank the members of Guillaume Rabusseau's group Mila, for the insightful discussions during the develop- ment of this project. Pascal Tikeng grateful Mahta Ramezanian- Sékou-Oumar Kaba, and Mohammad Pezeshki for helpful conversations in the early stages of this work. The authors acknowledge the material support of NVIDIA in the form of computational resources. Pas- cal Tikeng Notsawo acknowledges the support from the Canada Excellence Research Chairs (CHRC) program. Guil- laume Rabusseau acknowledges the support of the CIFAR AI Chair program. Guillaume Dumas supported by the Institute for Data Valorization, Montreal and the Canada First Research Excellence Fund (IVADO: CF00137433), the Fonds de Recherche du Québec (FRQ: 285289), the Natural Sciences and Engineering Research Council of Canada (NSERC; DGECR-2023-00N0 and the Canadian Institute for Health Research (CIHR 192031; SCALE).
Barak, B., Edelman, B. L., Goel, S., Kakade, S., Malach, E., and Zhang, C. Hidden progress in deep learning: Sgd learns parities near the computational limit. Neural Information Processing Systems, 2022. doi: 10.48550/ arXiv.2207.08799,
Beck, A., Levi, N., and Bar-Sinai, Y. Grokking at the edge of linear separability, 2024. URL https: //axiv. hrg/abs/24100 04489.
Bereska, L. and Gavves, E. Mechanistic interpretability for ai safety - review, 2024. URL URL.httss:/arrxix org/ abs 000s/2404.11 1.440.1148.2
Biehl, M. and Schwarze, H. Learning by on-line gradi- ent descent. Journal of Physies A: Mathematical and General, 28(3):643, feb 1995. doi: 10.1088/0305-44700 28/3/018. URL https: /
Impact Statement
This paper aims to advance the field of Machine Learning by improving the understanding of grokking in neural networks. While the focus on theoretical contributions, our work has practical implications for optimizing model training. We acknowledge the potential ethical considerations of Al technologies, but do not feel any specific issues need to be highlighted here at this time.
Boyd, S. and Vandenberghe, L. Conve Optimization. Cam- bridge University Press, 2004.
Bricken, T., Templeton, A., Batson, J., Chen, B., Jermyn, A., Conerly, T.. Turner, N., Anil, C., Denison, C., Askell, A., Lasenby, R., Wu, Y., Kravec, S., Schiefer, N., Maxwell, T., Joseph, N., Hatfield- Doddds Z., Tamkin, A., Nguyen, K., McLean, B., Burke, J. E., Hume, T., Carter, S., Henighan, T., and Olah, C. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. https:/transformer- circuits. wut.:p/ooooomoooommmaiicicccocc
References
Abramov, R., Steinbauer, F., and Kasneci, G. Grokking in the wild: Data augmentation for real-world multi-hop reasoning with transformers, 2025. URL https:// arxiv. or.oorg///a/oo abs /aabs//2504 20752.
htmm.
Candes, E. and Recht, B. Exact matrix completion via convex optimization. Communications of the ACM, 55 (6):111-119, 2012.
Aharon, M., Elad, M., and Bruckstein, A. K-svd: An algo- rithm for designing overcomplete dictionaries for sparse representation. IEEE Transactions on Signal Process- ing, 54(11):4311-422, 2006. doi: 10.1109/TSP.2206. 881199.
Candès, J. and Tao, T. The power of convex relaxation: Near-optimal matrix completion. JEEE transactions on information theory, 2655):2053-2000, 2010.
Arora, S., Cohen, N., and Hazan, E. On the optimization of deep networks: Implicit acceleration by overparame- terization, 2018. URL https: /atasi/.////i////a org/ abs/ 1802.06500
Candes, E. J., Romberg, J. K., and Tao, T. Stable signal recovery from incomplete and inaccurate measurements. Communications on Pure and Applied Mathematics: Journal Issued by the Courant Institute of Mathematical Sciences, 5908):207-11223 2006.
Arora, S., Cohen, N., Hu, W., and Luo, Y. Implicit regularization in deep matrix factorization. CoRR, abs/1905.13655, 2019. URL httht:///a/iiii
Chandrasekaran, V., Recht, B., Parrilo, P. A., and Will- sky, A. S. The convex geometry of linear inverse prob- lems. Foundations of Computational Mathematics, 12 (6):05- October 2012. ISSN 1615-3383. doi: 10. 10.100774102080012-993---- URL doi. org//10.1007/11000------1-1111
Aubin, B., Maillard, A., Barbier, J., Krzakala, F., Macris, N., and Zdeborová, L. The committee machine: Com- putational statistical gaps learning a two-layers
doi.
10

---

Grokking Beyond the Euclidean Norm of Model Parameters
Charton, F. Learning the greatest common divisor: plaining transformer predictions, 2024. URL https:
20132. URL https://onlleeeell///
wi wiley,
com/ com/ /atp///b////0/ htt.b////.1100/ppppp 20132.
//arxiv. orgg///gaa abs/ 15594.
Donoho, D. L. and Elad, M. Optimally sparse repre- sentation in general (nonorthogonal) dictionaries via &#x21 Bisup minimization. of the National Academy of Sciences, 100(5):2197- 2202, 2003, doi: 10. 10.1073/pnas 00.1073/pnas.443884770000 0437847100 URL
Chatterjee, S. Convergence of gradient descent for deep neural networks, 2022. URL https://arxi org/ abs/ 2200 64 62.
Proceedings
Chen, Y. Bhojanapalli, S., Sanghavi, S., and Ward, R. Co- herent matrix completion. In Xing, P. and Jebara, T. (eds.), Proceedings of the 31t International Conference on Machine Learning, volume 32 of Proceedings of Ma- chine Learning Research, pp. 674-682, Bejing, China, 22- 24 Jun 2014. PMLR. URL https: //ppporeeedngs... / vv32/
https:/ http://www..a..s.. wwt.ooog/oooooooooooooooooo /oi/ aat//aa///aa/ 10733/ pnas. 0437847100.
Doshi, D., Das, A., He, T., and Gromov, A. To grok or not to grok: Disentangling generalization and memorization on comrupted algorithmic datasets, 2024a. URL https: / //aaxivv.oo. org/ abs 2210 0.13.11061.1...
chenc14 httpa//s/v/cccccccccccccccc html.
Chughtai, B., Chan, L., and Nanda, N. toy model of universality: Reverse engineering how networks learn group operations. International Conference Machine Learning, 2023. doi: 10.485500ar21v.2332.0555
Doshi, D., He, T., Das. A., and Gromov, A. Grokking modular polynomials, 2024b. URL htttt:t///aaiiiii org/ abs /2406 03495.
Clauw, K., Stramaglia, S., and Marinazzo, D. Information- theoretic progress measures reveal grokking is emer- gent phase transition, 2024. URL https://// /arxiv. org/abs/ ataaab//4... 2408. 08944.
Engel, A. and Broeck, C. L. V. d. Statistical Mechanics of Learning. Cambridge University Press, USA, 2001. ISBN 0521773075.
Fan, S., Pascanu, R., and Jaggi. M. Deep grokking: Would deep neural networks generalize better?, 2024. URL https: www.xxxxv..rgaa///aoee
Cunningham, H., Ewart, A., Riggs, L., Huben, R., and Sharkey, L. Sparse autoencoders find highly interpretable features in language models, 2023. URL https:// arxiv. org/ abs 00.3090080000
19454.
Fazel, M., Hindi, H., and Boyd, S. A rank minimization heuristic with application to minimum order system ap- proximation. In Proceedings of the 2001 American Con- trol Conference. (Cat. No. No3014H37148), volume 6, pp. 4734 vol.6, 2001. doi: 10.1109/ACC.2201.45777
Czarnecki, W. M., Osindero, S., Jaderberg, M., Świrszez, G., and Pascanu, R. Sobolev training for neural net- works, 2017. URL 7706.04853
abs/
Fazel, M., Hindi, H., and Boyd, S. Rank minimization and applications in system theory. In Proceedings of the 2004 American Control Conference, volume 4, pp. 3273-3278 vol.4, 2004. doi: 10.23919/ACC.2004 1384522
Daubechies, I., Defrise, M., and Mol, C. D. An iterative thresholding algorithm for linear inverse problems with sparsity constraint, 2003. URL https://aaxix: org/ org/abs/neen babb//aath/33775
Foucart, S. and Rauhut, H. A Mathematical Introduction Compressive Sensing. Birkhäuser Basel, 2013. ISBN 08 0817649476.
Davies, X., Langosco, L., and Krueger, D. Unifying grokking and double descent, 2023. URL https:
//arxiv.org/ abs /2303 06173.
Gidel, G., Bach, F. R., and Lacoste-Julien, S. Implicit regularization of discrete gradient dynamics in deep linear neural networks. CoRR, abs/1904. 13262, 2019. URL httppp/aaaxiivoorrrbb///9..
DeMoss, B., Sapora, S., Foerster, J., Hawes, N., and Posner, The complexity dynamics of grokking, 2024. URL
https: /awaiiiiv..///..///aaaaa//aaaava org/ abs 2412.411222 09810.
Donoho, D. Compressed sensing. IEEE Transactions on Information Theory, 52(4): 1306, 2006a. doi: 10. 1109/TT.2200..775522.
Gissin, D., Shalev-Shwartz, S., and Daniely, A. The im- plicit bias of depth: How incremental learning drives generalization. CoRR, abs/1909. 2019. URL http: http:///acxiv.orr/abb////
Donoho, D. L. For most large underdetermined sys- tems of linear equations the minimal 4y-nrrm so- lution is also the sparsest solution. Communi- cations on Pure and Applied Mathematics, 59(6): 797-829, 2006b. doi: http:://oi..rr//0..
Goldt, S., Advani, M. S., Saxe, A. M., Krzakala, F,, and Zde- borová, L. Dynamics stochastic gradient descent for setup*.
1002/epa. 1002/////....
two-layer neural networks in the

---

Grokking Beyond the Euclidean Norm of Model Parameters
Journal of Statistical Mechanics: Theory and Experi- ment, 2020(12): December 2020. ISSN 1742- 5468. doi: 10. 10.1088/17422556664bbbo66 URL http:
Langley, Crafting papers on machine learning. In Langley, P. (ed.), Proceedings of the 17th International Conference Machine Learning (ICML 2000). pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
Golechha, S. Progress measures for grokking on real- world tasks, 2024. URL https://arxiv.or
Lee, J., Kang, B. G., Kim, K., and Lee, K. M. Grokfast: Accelerated grokking by amplifying slow gradients, 2024, URL https: arxiv. org/ / aaa//ba/// 2205/2455..0 20233.
abs/2405. 12755.
Gromov, A. Grokking modular arithmetic. arXiv preprint arXiv: Avxiv-2301.026799 2023.
Levi, N., Beck, A., and Bar-Sinai, Y. Grokking in linear estimators a solvable model that groks without under- standing. International Conference on Learning Repre- sentations, 2024.
Gunasckar, S., Woodworth, B. E., Bhojanapalli, S., Neyshabur, B., and Srebro, N. Implicit regularization in matrix factorization. Advances in neural information processing systems, 30, 2017.
Li, J. Mechanistic interpretability binary and ternary transformers, 2024. URL https://arxiv/
Abs/2405.17703.
Han, S., Pool, J., Tran, J., and Dally, W. J. Learning both weights and connections for efficient neural networks. CoRR, 0obs/1506.02622 2015, URL http::///// //axiv.. org/abs/ 606.006666 02626.
Li, Z.., Luo, Y., and Lyu, K. Towards resolving the implicit bias gradient descent for matrix factorization: Greedy low-rank learning. CoRR, abx/2012.098333 2020. URL https: httpss//aaxiv.... http//xxiivv.////// http://aaaii..o..a//////.///
09839.
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. Lora: Low-rank adaptation of large language models, 2021. URL https::t///... aattps://axt /arxiv. / 2106 09685.
Liu, C., Zhu, I.., and Belkin, M. Loss landscapes and optimization in aver-parameterized non-linear systems and neural networks, 2021. URL https: //at::v///xvv... org/ abs /aa//ab//20000 /2003. 00307.
Hu, M. Y., Chen, A., Saphra, N., and Cho, Latent state models of training dynamics, 2024. URL https: //arxivv org/ abs/ 2308 09543.
Liu, Z., Michaud, E. J., and Tegmark, M. Omnigrok: Grokking beyond algorithmic data. In The Eleventh International Conference on Learning Representations, 2023a. URL https: Setttt:://opeereeee /pwp//peeeeeeeeeeee/p/ee net forum? i id=2Di Ho 0Wa
Huang, Y.. Hu, S., Han, X., Liu, Z., and Sun, M. Unified view of grokking, double descent and emergent abilities: A perspective from circuits competition, 2024. URL https:/ aaaaiaiiiv.v//.. www...www................ooo orgo absg//////// 2402. 15175.
Liu, Z.,, Zhong, 7., and Tegmark, M. Grokking as compres- sion: nonlinear complexity perspective. arXiv preprint arXiv: 2310.05918, 2023b.
Huggins, J. H., Campbell, T., and Broderick, T. Coresets for scalable bayesian logistic regression, 2017. URL
Lu, L., Meng. X., Mao, Z., and Karniadakis, G. E. Deepxde: A deep learning library for solving differential equations. SIAM Review, 63(1):208 -228 Jamuary 2021. ISSN 1095- 7200. doi: 10.1137/19m1274007. URL http://dx: doi org 10. 1137 1137/19112111071711
https: /aass///arxiiie //aaariivvv///.. org/ absaabss/555 0423.
Humayun, A. I., Balestriero, R., and Baraniuk, R. Deep networks always grok and here is why, 2024. URL. atttss://arxx orgg abs/ 2202. 15555.
Lucic, M., Faulkner, M., Krause, A., and Feldman, D. Train- ing gaussian mixture models at scale via coresets. Jour- nal of Machine Learning Research, 18(160):1-25, 2018. URL http: /ttt:://jmrr../ http:ilttt.paaaaeemeeeeeeeeeeeee wwtpp//ppee/////1------. html.
Karimi, H., Nutini, J., and Schmidt, M. Linear conver- methods under the tolyak-tojasiewiiz condition, 2020. URL. https:
gence gradient and
//arxiv. orgo abs 04636.
Kumar, T., Bordelon, B., Gershman, J., and Pehlevan, C. Grokking the transition from lazy to rich training dynamics. arXiv preprint arXiv: 23 23/066 10, 2023.
Lyu, K., Jin, J., Li, Z., Du, S. S., Lee, J. D., and Hu, W. Di- chotomy of early and late phase implicit biases can prov- ably induce grokking. arXiv preprint arXiv: 1.181, 2023.
Kumar, T., Bordelon, B., Pehlevan, C., Murthy, V. N. and Gershman, S. J. Do mice grok? glimpses of hidden progress during overtraining in sensory cortex, 2024.
Mairal, J., Bach, F., Ponce, J., and Sapiro, G. Online learn- ing for matrix factorization and sparse coding, 2010. URL https: /tt///axiii// http//avvrrgg.baaaeeeee//e////ee 09000000000000 0050.
URL wwtp:/xwwwog/////////xo 03541.
12

---

Grokking Beyond the Euclidean Norm of Model Parameters
Mairal, J., Bach, F. R., and Ponce, J. Sparse modeling for image and vision processing. CoRR, abs/1411.3230, 2014. URL bestp://arxiv/
Natarajan, B. K. Sparse approximate solutions linear systems. SIAM Journal on Computing, 24(2):227-234, 1995. doi: 10. 10.11377509955339992440666 URL https: //doi. /ttp://.....11/00077399992220000
http//bss//1111 3230.
Mallinar, N., Beaglchole, D., Zhu, L.,, Radhakrishnan, A., Pandit, P., and Belkin, M. Emergence in non-neural mod- els: grokking modular arithmetic via average gradient outer product, 2024. URL https::////. wwwpiaawwv.og..............
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. arXiv preprint arXiv: 2306. 13253, 2023.
Merrill, W., Tsilivis, N., and Shukla, A. A tale of two circuits: Grokking as competition of sparse and dense subnetworks. arXiv preprint arXiv: 2303. 1187,, 2023.
Olshausen, B. A. and Field. D. J. Emergence of simple-cell receptive field properties by learning sparse code for natural images. Nature, 381:607-609, 1996. URL https: // /ttppaapp..///// www.steesmannnnneeeeeeee org/ /
Miller, J., Gleeson, P., O' D'Neill, C., Bui, T., and Levi, N. Measuring sharpness in grokking, 2024a. URL https:
// rxiv. org/ abs/ 2402. 08946.
Park, Y., Kim, M., and Kim, Y. Acceleration of grokking in learning arithmetic operations via kolmogorov-amool representation, 2024. URL abs/ a 16658.
Miller, J., O'Neill, C., and Bui, T. Grokking beyond neural networks: An empirical exploration with model complex- ity, 2024b. URL https:tta..... : httppp:///aaxii.//eeaeeeessassss arwia. org/ www.ars//////////////..//..ae/a/ abs/abb/..1.... 17247.
Power, A., Burda, Y., Edwards, H., Babuschkin, I., and Misra, V. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv; Arxiv- 2201.02177, 2022.
Minegishi, G., Iwasawa, Y., and Matsuo, Y. Bridging lottery ticket and grokking: Understanding grokking from inner structure of networks, 2025. URL pattps://arxx org/abs 2310. 19470.
Prieto, L.,, Barsbey, M., Mediano, P. A. M., and Birdal, T. Grokking the edge of numerical stability, 2025. URL https://arxien org/ ats//bbs//2200 04697.
Mohamadi, M. A., Li, Z., Wu, L., and Sutherland, D. Why do you grok? a thcoretical analysis of grokking modular addition, 2024. URL https: /ttxivaavvvv//////////ee.eawwwvl org/ abs/ 2407. 12332.
Raissi, M., Perdikaris, P., and Karniadakis, G. Physics- informed neural networks: A deep learning frame- work for solving forward and inverse problems involv- ing nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, 2019. ISSN
Morwani, D., Edelman, B. L., Oncescu, C.-A., Zhao, R., and Kakade, S. Feature emergence via margin maximization: studies in algebraic tasks, 2024. URL https:// arxiv. www.vv.rg///ooo abs 2311.01100000 07568.
0021-9991. doi: https://biiioggg///..00ddd..pppp
Murty, S., Sharma, P., Andreas, J., and Manning, C. D. Characterizing intrinsic compositionality in transformers with tree projections, 2022. URL https: //arxi...vvvv/////////aa//aaa org/abs// 2211. 01288.
045. URL wtt.D//tpp///ww..ccceeeeeeeeeeed sctennee/ariicee///////219999999
Rauhut, H. Compressive Sensing and Structured Ran- dom Matrices, pp. 1-92. De Gruyter, Berlin, New York, 2010. ISBN 9783110226157.
Murty, S., Sharma, P., Andreas, J., and Manning, C. D. Grokking hierarchical structure in vanilla transform- ers, 2023. URL. https:ttss.... http///axvvv//ee///e/eee/eaaa wwwpxiivvooorra.ro 22055. 18741.
doi: doi:10. 1515/9783110226157.11 URL ptttps://oi..oo 10. 10.1515/978311022555
Mustafa, N. and Burkholz, R. Dynamic rescaling for train- ing gnns. In Globerson, A., Mackey, L.. Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Advances in Neural Information Processing Systems, vol- Curran Associates, Inc.,
Razin, N. and Cohen, N. Implicit regularization in deep learning may not be explainable by norms. CoRR, abs/2005.06398, 2020. URL /tttp:://arxiiiii abs/2005.06333
ume 37, pp. 103370-125393, 2024.
Recht, B., Fazel, M., and Parrilo, P. A. Guaranteed minimum-rank solutions of linear matrix equations via nuclear norm minimization. SIAM Review, 52(3): 471-501, January 2010. ISSN 1095-7200. doi: 10. 1137/070697835. URL http: /dx doi..ii.... dow..rgg10........ 1137/070697835.
Nanda, N., Chan, L., Lieberum, T., Smith, J., and Stein- hardt, J. Progress measures for grokking via mechanistic interpretability. International Conference On Learning Representations, 2023. doi: 10.485500/rXivv2330
13

---

Grokking Beyond the Euclidean Norm of Model Parameters
Rockafellar, R. T. Convex analysis. Princeton Mathematical Series. Princeton University Press, Princeton, J.,,, 1970.
Varma, V., Shah, R., Kenton, Z., Kramár, J., and Kumar, R. Explaining grokking through circuit efficiency. arXiv preprint arXiv: 2309.02399 2023.
Rubin, N., Seroussi, I., and Ringel, Z. Grokking as first order phase transition in two layer networks, 2024. URL https: /ttp///axvvvv/vv///vv///////.avv org/ abs 0...0.0.7788
Vavskevivcius, T., Kanade, V., and Rebeschini, P. Implicit regularization for optimal sparse recovery, 2019. URL https: /ttp//aaxivvv/v. wwtpxiiv.orr//eoo aaaaaass19905 05122.
Saad, D. and Solla, S. Learning with noise and regu- larizers in multilayer neural networks. In Mozer, M., Jordan, M., and Petsche, T. (eds.), Advances in Neural Information Processing Systems, volume 9. MIT Press, 1996. URL cc cc/paper_ /atpppeeffieee/ppppp/////////L// e e1d5bel c7 (aa.eeeccc7744556666666333331130 pdf.
Wang. B., Yue, X., Su, Y., and Sun, H. Grokked trans- formers implicit reasoners: A mechanistic journey to the edge of generalization, 2024. URL https: / //arxiv. wwwrxii.orr/abs/2405....
Wedin, P.-A. Perturbation bounds in connection with singular value decomposition. BIT Numerical Math- ematics, 12(1))99-111, 1972. ISSN 1572-9125. doi: 10. 10.1007/BF01932677.. URL http:: /doi org/10. 1007 #007/BF0193665
Saad, D. and Solla, S. A. On-line learning in soft committee machines. Phys. Rev. E, 52:4225-4243. Oct 1995a. doi: 10.1 10.1103/PhysRReF555242255 ggs. 10. http://///1..11///P1e1/. 11C3//hhyPP.........m 52, 4225.
URL https: /link.
Xu, Z.. Wang, Y., Frei, S., Vardi, G., and Hu, W. Benign overfitting and grokking in relu networks for xor cluster data, 2023. httht::: : /arxiv..vvv. org/ abs 2310. 02541.
Saad, D. and Solla, S. A. Exact solution for on-line learn- ing in multilayer neural networks. Phys. Rev. Lett., 74: May 1995b. doi: 10. 10.1103/hhssRvvLett. 74.4337. URL http:: wtpp//innknn..........mmmpmms..m org/doi/10. 1103/PhysRevLete. 4337.
Xu, Z., Ni, 7.,, Wang, Y., and Hu, W. Let me grok for you: Accelerating grokking via embedding transfer from a weaker model, 2025. URL https: http///axxvvorrr/e abs abs/2504. 13292.
Song, K., Tan, Z., Zou, B., Ma, H., and Huang, W. Unveil- ing the dynamics of information interplay in supervised learning. 2024. URL https: /ttp:://axxvvvvii wwwpxii..oor//ooooooooo abs/ 2406. 03999.
Yunis, D., Patel, K. K., Wheeler, S., Savarese, P., Vardi, G., Liveseu, K., Maire, M., and Walter, M. R. Ap- proaching deep learning through the spectral dynamics weights, 2024. URL https: wwwaaiii.oo//ba........eee 2408 11804.
Stander, D., Yu, Q.. Fan, H., and Biderman, S. Grokking group multiplication with cosets, 2024. URL https: //arxiv. org/ abs/2312 06581.
Zhou, X., Pi, R., Zhang, W., Lin, Y., Chen, Z., and Zhang, T. Probabilistic bilevel coreset selection. Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), Proceedings of the 39th International Confer- ence Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 27287-27302 PMLR, 17-23 Jul 2022. URL https::: pesttts://proeeeeee http///pcceeennn.eeec .mlr. press/ v1 62. zhou22h. html.
Thilak, V., Littwin, E., Zhai, S., Saremi, O., Paiss, R., and Susskind, J. The slingshot mechanism: An em- pirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint arXiv: Sxiv-2206.04817, 2022.
Tibshirani, R. Machine learning 10-725/36-725, con- vex optimization: Spring 2015, lecture 8: February 9. https://wwet stat. cmu wewma.eduu-eee hammmedddrrryataaa. convexopt 2015.
Zhou, Z., Zhang, Y., and Xu, Z.-Q.
J. A rationale from frequency perspective for grokking in training neural network, 2024. URL htttt:t////xi..ooo /ttp//a.........aa...aa.......aa 2405.17473
Tsang, I. W., Kwok, J. T., and Cheung, P.-M. Core vec- for machines: Fast svm training on very large data sets. Journal of Machine Learning Research, 6(13):363- 392, 2005. URL http: http:://mmrorr////// /ttp:///////or/////er//////// papers/v6/ tsang05a.htmn
Zhu, X., Fu, Y., Zhou, B., and Lin, Z. Critical data size of language models from a grokking perspective, 2024. URL https::/p///. ata/aaiivvvvr.... aww./b/40000 10463.
v Zunkovive, B. and llievski, E. Grokking phase tran- sitions in learning local rules with gradient descent. Journal of Machine Learning Research, 25(199):1 httttpp://mm....ooppppprr///////
Vandenberghe, L. and Boyd, S. Semidefinite programming.
SIAM Review, 38(1):49-9))) 95, 1996. doi: 10. 10.1137/1038003.
52, 2024. httht://// /tt.r///.///////////////// waml..org/paan orgg/papers/v25/ 22-1228. html.
URL https: /doii wttp:///oooooo.....oooo/.o /1t. /1038003.
14

---

Grokking Beyond the Euclidean Norm of Model Parameters
A. Related Works
16
A.1. Grokking
16
A.2. Model Distillation, Sparse Dictionary Learning, Sparse Auto-Encoder, Coreset Selection
18
A.3. Finite Field
18
B. Notations
19
C. Proofs
19
C.1. Proof of Theorem 2.1
19
C.2. Main Theorem under Polyak-Lojjsiewiizz Inequality
30
C.3. Preservation of the Chatterjec-LLojssieiizz Inequality under Perturbation
31
C.4. Proof of Theorem 3.1
34
C.5. Proof of Theorem 3.3
39
C.6. Proof of Theorem 3.4
40
C.7. Proof of Theorem 3.5
43
C.8. Proof of Theorem 3.6
44
D. Other Iterative Method for and l, Minimization
45
D.1. Sparse Recovery
45
D.2. Low Rank Matrix Factorization
47
E. Grokking Without Understanding
48
Implicit Bias the Depth
49
F.I. Deep Sparse Recovery
49
F.2 Deep Matrix Factorization
50
G. Amplifying Grokking through Data Selection
52
G.1. Sparse Recovery
52
G.2. Matrix Factorization
56
H. Additional Experiments
60
H.1. Sparse Recovery
60
H.2. Matrix Factorization
61
H.3. General Setting
61
H.3.1. Algorithmic dataset
64
H.3.2. Non linear Teacher-Student
65
H.3.3. Domain Specific Regularization
65
H.3.4. Image Classification
67
15

---

Grokking Beyond the Euclidean Norm of Model Parameters
A. Related Works
A.1. Grokking
Universality Liu et al. (2023a) took the first step in demonstrating the universality of grokking by inducing it on non- algorithmic data, notably on image classification, sentiment analysis, and molecule property prediction. On the same line, Gromov (2023); vZunkovive & Ilievski (2024); Levi et al. (2024) show that grokking can occur for (analytically solvable toy) models. Barak et al. (2022) observed grokking on the binary sparse parity problem, Charton (2024) on greatest common divisor (with Transformer), Doshi et al. (2024b) modular polynomials, Lyu et al. (2023) on matrix completion and sparse linear predictors, Beck et al. (2024) and Levi et al. (2024) on binary logistic classification, Xu et al. (2023) on XOR-cluster training data. Miller et al. (2024b) show that grokking occurs beyond neural networks-in models like Gaussian process (regression and classification), linear regression, and Bayesian neural networks- driven by a trade-off between model complexity and error, not just feature learning or weight decay. Wang et al. (2024) and Abramov et al. (2025) show that grokking enables Transformers to develop reasoning abilities, emerging only after extended training, whether on synthetic comparison/composition tasks or real-world multi-hop reasoning augmented with inferred facts. Li (2024) observe that both binary and ternary Transformer exhibit grokking modular addition under weight decay regularization, and do not do so without weight decay. Mallinar et al. (2024) show that grokking occurs beyond neural networks and gradient descent, as kernel- based Recursive Feature Machines trained with Average Gradient Outer Product also exhibit delayed generalization through emergent structured features. Kumar et al. (2024) find that mice show grokking-like behavior, i.e., even after behavioral performance plateaus, neural representations through margin maximization. Mustafa & Burkholz (2024) investigate grokking- like phenomena in Graph Neural Networks, particularly focusing on tasks involving heterophilic and homophilic graphs
sensory cortex continue evolving, improving generalization
Interpretability and Explainability Since the concept of "grokking" emerged, numerous theories have been proposed to explain the phenomenon. Liu et al. (2023b) present representation learning the main underlying factor of the existence of generalizing solutions, supporting Power et al. (2022)'s 's preliminary observations. Exploring the direction that separates generalization from memorization solutions (i.e., progressions measures for generalization), recent studies shed light on various factors, such as neuron activity (Nanda et al., 2023), weight norm of the model parameters (Liu et al., 2023a), sparsity (Merrill et al., 2023), time scales of pattern formation (Davies et al., 2023), Fourier gap (Barak et al., 2022), last layer norm (Thilak et al., 2022), fast vs low-frequency components (Zhou et al., 2024), mutual information ratio of neural representation (Song et al., 2024), phase transition (Clauw et al., 2024; Rubin et al., 2024), rank minimization (Yunis et al., 2024), rise and fall in model complexity (DeMoss et al., 2024; Humayun et al., 2024), etc. Also, weight decay and weight norm decrease (Liu et al., 2023a; Varma et al., 2023), a gradual process facilitated by optimization (Nanda et al., 2023; Merrill et al., 2023; Barak et al., 2022; Davies et al., 2023; Notsawo et al., 2023), the instability induced by the Adam optimizer (Thilak et al., 2022), are among the factors that has been use to explain why the transition from memorization to generalization. According to Varma et al. (2023), grokking arises from a generalizing circuit being favored over a memorizing circuit. Miller et al. (2024a) introduce a method to quantify the sharpness of the grokking transition in neural networks. Doshi et al. (2024a) show that in the presence of label corruption, grokking can still occur, with networks first generalizing before later unlearning memorized noise, highlighting a separation between generalization and memorization phases. Zhu et al. (2024) show that grokking in language models emerges when training data exceeds critical size, with larger models requiring more data to generalize beyond memorization. Huang et al. (2024) propose a unified framework explaining grokking. double descent, and emergent abilities through the lens of competition between memorization and generalization circuits. Chughtai et al. (2023) and Stander et al. (2024) reverse engineer (small) neural networks learning finite group composition via mathematical representation theory. Golechha (2024) show that lg norm alone can not explain grokking, as it occurs even outside the typical "goldilocks zone"(Liuu et al., 2023a), and propose alternative measures like activation sparsity and weight entropy that better track generalization dynamics. Prieto al. (2025) explain grokking as a instability from floating. Moating-point errors that halts learning after memorization.
result of Softmax Collapse-numend
Predictability To the question of whether grokking can be predicted, Notsawo et al. (2023) answered in the affirmative, proposing predictable measure the spectral signature of the training loss in the early optimization phases of the model training. Hu et al. (2024) models neural network training uncover discrete training phases. Applied to grokking tasks, this reveals distinct latent states corresponding to memorization and generalization, with "detour" states explaining delayed generalization. The approach offers a structured way to analyze and predict grokking by linking it to transitions between these latent training states. Murty et al. (2023) explore how vanilla Transformers eventually grok hierarchical rules after extended training, and find that weight norm and attention sparsity
Hidden Markov Model (HMM) over weight statistics to
16

---

Grokking Beyond the Euclidean Norm of Model Parameters
are inconsistent and often unreliable for forecasting generalization. They introduce "tree-structuuednness (Murty et al., 2022) metric that significantly outperforms weight norm in predicting whether and when model will grok. Miller et al. (2024b) argue that grokking may be possible in any model where the solution search is guided by complexity and error. This notion of complexity is related to what we refer to as regularization in our work. Fan et al. (2024) show that deeper MLPs on MNIST grok more often and in stages, and that weight norm poorly predicts generalization compared to internal feature rank dynamics.
Stochasticity Our work also contradicts the hypothesis put forward when grokking was first observed, namely that grokking may be due stochastieity or an anomaly in the optimization (Power al., 2022; Thilak al., 2022). For sparse recovery and matrix factorization, the optimization algorithms we use are all deterministic (up to initialization).
Feature Learning Kumar et al. (2023) show that grokking can emerge without weight decay, as result of a delayed transition from lazy (fixed-feature) to rich (feature-learning) dynamics during training. In contrast, Lyu et al. (2023) explain grokking through implicit bias, showing that large initialization and small weight decay induce a shift from kernel-like to margin-maximizing behavior, highlighting different but complementary mechanisms behind grokking. Xu et al. (2023) provide a theoretical example of grokking in ReLU networks by analyzing the feature learning process under GD, and show that test generalization can emerge well after perfect (and initially non-generalizing) memorization of noisy XOR-cluster training data. Morwani et al. (2024) demonstrate that neural networks trained on tasks like modular addition and sparse parity naturally converge to solutions based on Fourier features and group- toup-thee representations (Nanda et al., 2023; Chughtai et al., 2023). This behavior arises from the principle of margin maximization, even without explicit regularization, providing a theoretical explanation for grokking as a delayed transition from memorization to structured feature learning.
Large Initialization and lg Regularization A direct work related to ours on grokking is Lyu et al. (2023), which analyzes the emergence of grokking under lg regularization and large-scale initialization in homogeneous models. Their results show that under such conditions, grokking can occur as a delayed generalization phenomenon. However, their analysis is limited to l weight decay and does not generalize to other regularization schemes. In contrast, our work extends beyond fg. demonstrating that grokking also arises under & and nuclear norm (f.) regularization, even without large-scale initialization. In the theoretical settings we study-such as sparse recovery and low-rank matrix factorization-66, regularization alone fails to induce grokking. Instead, it produces a sharp drop in generalization error during training that does not correspond to true generalization. We refer to this failure mode as grokking without understanding, which we attribute to the violation of key assumptions in Lyu et al. (2023) (notably Assumption 3.2). Our results further show that grokking may even be necessary in practice: for instance, when targeting low-rank solutions via L. regularization, it is preferable to use the smallest feasible regularization strength and train well beyond the point of overfitting to achieve generalization. Levi et al. (2024) work on classification settings and show that the sharp increase in generalization accuracy may not imply a transition from "memorization" " to "understanding" but can be an artifact of the accuracy measure. This aligns with the "grokking without understanding" problem we observe in sparse recovery and low-rank matrix factorization.
Accelerating Grokking Park et al. (2024) accelerate grokking by using data augmentation for commutative operations and transfer learning across model components aligned with the Kolmogorov-Amold representation. Lee et al. (2024) accelerate grokking by amplifying slow gradient components, reducing training time across tasks. Xu et al. (2025) accelerate grokking by initializing larger models with embeddings from smaller trained models, eliminating delayed generalization. Minegishi et al. (2025) show that grokking can be accelerated by identifying "grokking tickets"'' -sprrs subnetworks that emerge after generalization-ann retraining them, which leads to much faster generalization than the original dense model. (Prieto et al., 2025) accelerate grokking using a stable softmax variant and a gradient projection method to prevent collapse and enable faster generalization without regularization.
Sparsity and Low-Rankness Barak et al. (2022) observed grokking on the binary sparse parity problem, and Merrill et al. (2023) shows that two subnetworks compete during training on such task: a dense (memorization) subnetwork and sparse (generalization) subnetwork. Since very sparse network that generalizes the sparse parity data can be built (Merrill et al., 2023), we conjecture that it is this sparsity that gives the models trained on this task their grokking nature, as well as the algorithmic dataset (Power et al., 2022), but leave further investigation for future work. To the best of our knowledge, we are the first to formally study grokking in the context of sparse recovery and low-rank matrix factorization (the shallow case). Lyu et al. (2023) show that low-rank matrix completion problems exhibit grokking with large initialization. But we
17

---

Grokking Beyond the Euclidean Norm of Model Parameters
prove that even on such simple model, we do not need way decay and large initialization observe grokking, but just f1/+ regularization.
A.2. Model Distillation, Sparse Dictionary Learning, Sparse Auto-Encoder, and Coreset Selection
Note that the parameterization b &a gives Fa (x) network with the weights of the first layer given by Ø E Rm*x (fixed) and those of the second layer given by a E R"" to be optimized. Compressed sensing thus corresponds somehow find a sparse version of model parameters using fewer data points. This could be strengthened by assuming that we have only the measures y* of the sparse representation a" of the signal and the measurement matrix M, but not the basis in which it is sparse, and the problem is to find both and a", i.. minimize ||al1 subject to F (M) y"Il2 Σ e and & E C. This is the case for matrix factorization, where the sparse basis (singular vector space) is jointly optimized along with the sought sparse coordinates a* (singular values). Here, C can be the set of orthonormal matrix (Φ & In), unit column norm ",, . 1), etc. It can be interesting see if l allows to recover the parameters with fewer samples than the matrix Φ
& TT x, i.e. Fx(X) X which is a two-layer linear neural
model distillation (teacher-student), where we aimed to
most used fg in deep learning. Looking for sparse a mean, we assume that only few components of the representation h(x) = + x E R" contribute to the output (x) hxx). This kind of method (sparse auto-encoder) is used a lot in mechanistic interpretability today to separate and interpret the features a pre-trained model has learned (Cunningham et al., 2023; Bricken et al., 2023), a promising direction for AI safety (Bereska & Gavves, 2024). Compressed sensing theory gives us an idea of how to select data to extract important features with the fewest possible examples: we need to select the samples that are most "incoherent" with the feature extractor (assumed fixed).
This formulation is also related to:
The coreset selection problem (Tsang et al., 2005; Huggins et al., 2017; Lucic et al., 2018), which consists of selecting small subset comprised of most informative training samples such that training on this subset can achieve comparable or even better performance with that on the full dataset (Zhou et al., 2022).
The sparse dictionary learning problem (Olshausen & Field, 1996; Aharon et al., 2006; Mairal et al., 2010; 2014), where we have a collection B e Rdx of N point in Rd and aim to write each of them as a combination of a few numbers of atoms of a dictionary & E Rdxm, that is B . [N], with ||o small. Let C {C |pdxn I|C.., ll2 S
Vi E [n]]. We want क = arg mingec N 2i=1 I(B T)) with 1(b, b) = minaeR"
N
fall3 + A A||a||i That is |
$ A
arg min
B
E + A
Σ
|A
AER.AE
N
The generalization error is [()) = Eb, t). The sparse coding problem 1(b, b) = minaer Rm ||b just the problem we study above.
+all2 + A||a|| is
We leave these points as a future direction for this work.
A.3. Finite Field
Compressed Sensing In compressed sensing and sparse recovery in F, = 7/p7 (for a prime integer p), incoherence between the measurement matrix M and the sparse basis t remains important but requires modular arithmetic. Recovery of a sparse signal b* from the measurements y* = M&b" is possible, provided that M and & are chosen to ensure low coherence in the modular structure. However, recovery algorithms are more complex due to the finite nature of the field and may involve different mathematical tools (e.g., algebraic techniques from coding theory) instead of standard optimization methods.
Matrix Factorization We can see A Ti X as representing a binary operation over F p. and the aim is learn this operation from a proportion of sample". This the vanilla setup where grokking was first observed (Power et al., 2022). Large-scale initialization can not explain the grokking phenomenon in such setup since it often observed with standard initialization. Many works also use mechanistic interpretability to find the algorithm deep neural networks use generalize
For example, A uu T with u = [O..... . 1] for the multiplication operation.
18

---

Grokking Beyond the Euclidean Norm of Model Parameters
in such a setup and propose progression measure for generalization (Nanda et al., 2023; Gromov, 2023). Still, none explain what makes such data particular. Liu et al. (2023b) propose the representation learning hypothesis but uses loss function that encourages such representation to emerge for standard deep neural networks (MLP and Transformer). Mohamadi et al. (2024) explain grokking in modular addition as transition from a kernel regime, where generalization fails, to a rich regime that captures the task's global algebraic structure- making grokking especially pronounced modular arithmetic. Investigating how our work can effectively answer why grokking is naturally observed in such a setup is a future direction for this work.
B. Notations
Our notation is standard.
We use the notation iid for "independently and identically distributed".. For A E R"**", the notation A iid N (H, means the entries of A are independently sampled from the normal distribution with mean μ and standard deviation a. This notation is valid for any other distribution used in this paper.
For two functions @, v Rgo + R, we write (()) (z)))) if there exist constants C 0 and 2g E Rgo such that [ó(z) S C' (z) for all z 2 zot d(z) O(v(z)) ((2) 2 Cz)(2) for all zp: and ((z) (vb(z) if both
x(2) = (x)(2) and Ø(z) = (2)(
We let e [ be the tt vector of the canonical basis of R" ek = 8g.Vl. The subscript (n) will be omitted when the context is clear.
(n)
max / (A) is the maximum (resp. minimum) singular value of a matrix A, with Amax / min (A) the corresponding cigenvalue
For a vector x E IR", ||x||o = Hi e [m], x; + 03/, ||x||p = (Ei (x Ap E (0, oo) and ||x||0 = max |x/|-
For a matrix A Rmxx, the schatten p-norm of A is | E,, (A))) where {A)(A)}: the set of singular value of A. For p = 1, this gives the trace/nuclear norm |A||. (.) = tr (VATA) A). The induced |x supp|i |xxx] We have ||A||1-1 maxje(m) E A |
p 9 norm of A is ||A|| ppp- s
(maximum absolute column sum). ||A||2-2 Alla = max (A) (operator norm, spectral norm, induced 2-norm) and || || 100-000 = max
Σi=1 |Ay (maximum absolute row sum).
O is Hadamard product. For A E Rmxn and B E Rmxx (A O B)i Bi, (0 m. 0 0S0 < <)
@ is the Kronecker product. For A E |Raxn and B E RPx8, (A & B)p A,.Brr.. (0 S < m. 0 <
0 < and 0 < w < q).
For A E Rmxx and B E Rex*, the Khatri-Rao product A * B E RRmxxn contains in each column i E [n] the matrix
A,, & B We have the formula A + B = (A & 1,) O (Im @ B).
For A E Rmxn and B E RMxx, the face-splitting product A o B E RuxXp contains in each row i E [m] the
matrix A @ B It can be seen as the row-wise Khatri-Rao product, and we have (A B) = (A T * B B)) ④ [1, 8 B).
(A 8
C. Proofs
C.1. Proof of Theorem 2.1
Let g: RP + |0, xo) be differentiable function, h: : IRP + [0, oo) be a subdifferenntiale (often convex) function and 3 >>>>>>>>>> 0. We want to minimize f == + .hh using gradient descent with a learning rate > O. The subgradient update rule for this problem is given by
x ((+1) = (0) aF(x!0)= (0 - (c(x())+3H(x())
3H(x())) v 0
(3)
19

---

Grokking Beyond the Euclidean Norm of Model Parameters
where F()) E af(x) G(x) + 330(x), with G(x)= = Vg(x) the gradient of g at x and H()) E ah(x) any subgradient of hxx) at We want show that, under certain conditions on and /3, this procedure will first converge to solution & that minimize 9 after 1 < oo training steps, then take additional At = o m3 training steps converge to solution x* that minimizes both f and h. Subdiffereniiability. convexity, and smoothness are defined below for completeness.
Definition C.1 (Subddifferentiability). A function yp: IRP + IR is said to be subdifferentiable at x E IR? if and only if there x) p-0xx
exists z E RP, <py))=>)))+)))>>)))) 2 (xx) + (y- denoted (x). When 0=(x) is the singleton set, we say that is differentiable at x 0x(() {V+((x))
v'y E RR. The set of all such z at is called the subdifferential of 4p at x and is
Definition C.2 (Convexity). A function y : RD + R is said to be convex if and only if o (tx + (1 - t)y) I ttp(x) + (1 ()((y) vx,y E RR, AI E (0, 1). If sp is subdifferentiable, this implies (yy) 2 (xx) x)) z VX, y E R2, Va E dp(x). If to is twice-differentiable, this implies Amnin (x)) 2 0 VX E RP.
Definition C.3 (Smoothness). Let Φ: Rp + R be a differentiable function and L > O. We say that is L-smooth if and only if V is L-Lipschitz continuous, i.e. (-) - V & L - x112 Vx,y E R?.
Lemma C.4. If function 9: : R + IR is L-smooth, then
p(y) S sx(x) (y x) Vx(()) + fly - v E IR, The converse is salse in general, and true for a corvex s.
I[Vx(x)11 S (x) VX E RD for non-negative sp The converse is false.
X(V E [- VX Ro for twice-differentiable sp The converse is true.
Vy is L/L-coceercive for a convex y, i.e. (Baillon-Haddad inequality). The converse is Irue.
(x yx 2 1IVA(y) Vv(x)) Vx,y RD
We use the following notations. For a non empty set θ c Rp and a vector x E Rp, we let dist(x, Q) == infyee ||x ex-yle yllg. We define f* x (x) and , arg min (x). Similarly we define g" and Og h² and On- We finally set hy inf xee, h(x), i.e. h infxERr hxx) s.t. gg. and 0% == arg min h(x).
In general, f * g" + 0hh. : In fact, we have f 2 g" + Bh.. This implies
f 2 g" + Bh + ff" S -g" - -8h"
= -0s -5" S fxx -g" -B"- g(x) 9" + +(x(x)) - h²))
-S'S
(4)
= 15(x) - 1'1 11(()= - -g" +B((xx) - 1(') (l(x)) g01 +/11(xx))) - -11
So g(x) - g" and h(x) + h- implies f(x) + f" But in general, it will not be possible to jointly have g(x(()) g" and h(x(0)) + h² by gradient descent. We will show that we first have y ((x(1)) - g", then h((x(0) hig
Assumption C.S. , n * 0.
This assumption assumes too little noise compared the signal. Even if initially 9(x(1) g" (i.e., memorization of training data with associated noise), as soon as generalization occurs, lg(x(()) - g"1 becomes proportional to the noise in the data. But under this assumption, when memorization is achieved, we can not have (x(()) + f" without lx(x()) + hgg
Lemma C.6. If assumption C.5 holds, then f g" + Shg: and hence 3(h(x) hg) < g(x) g" + +((((x)) hgg) <
f(x) - f Vx E RR.
Proof. Assume ,, nyne 0, + 0. Then there exists x; 2, such that g(x)) = g¹. This implies x²- E Gg. Under this assumption, f s(x]) + Bh(x)) = 9" + 3h(x)). Since xj E Gg h(x}) 2 hū so f g" Bhg. Conversely, we have g" + Bhg g(x) + 3h(x)) f(x) 2 fr for any x 0g c G, Therefore, f = g" + Bhg This implies
f(x) f =g(x= g" + 3((()) - -h1) 2 ((((xx) - -ig) since g(x) - g" 2 0.
Assumption C.7. Since , = arg min (g(x) - 9 gg),, we assume without loss of generality g" = 0.
If g is the loss function of an overparameterized neural network, then for 3 = 0, f = g converges to a solution x* close to the initialization x)) such that g(x*) ) = 0, under certain conditions on a and g (see, for example, Theorems C.25 and C.15). We will extend this result to the case / > 0.
Assumption C.8. We assume there exists x E R2 such that g(x) = g"
20

---

Grokking Beyond the Euclidean Norm of Model Parameters
We define the Shatterjee-Eojasienee (CL) constant of subdifferentiable function ş IRP + [0, 00) x RP with radius > and with respect to another function Φ : RP [0, 0o) as:
8
if oly) - OVy E B(x,r)
X
x,,, ==
infy ((xxrr) oly)#0
inf e 1
otherwise.
(ly)
(5)
00
if fy) - E r(x,r) otherwise.
V (ly)
up is differentiable
infye (xxx))
<(y)<0
and
:= where (xx, r) = (y E Rp ||x 1 y rr denote the closed Euclidean ball of radius r centered at x. inequality). For r >> 0, a nonnegative subdifferentiable function yp RP - [0, oo)
(6)
Definition C.9 bebaaterjee-Loojsieeeeee
is said to satisfy the r-CL inequality at x E IR if and only if 4Ø(x) < x,,,
Theorem C.10. Take any x( x
E Ro with g( := g)x + O. Assume that g satisfy the r-CL inequality at x
x()) for some
0, Le. 4g(0) < x x(g,x((),, r). There exist fmax > O, Ommx 0 and three constants C, C', C" 0 such that for
all a E (0, mmaxx) and 3 E (0, fmax). by defining the subgradient descent update (3) with and 3 starting at x()), the
following hold:
Fo any Eg (( there exists a step t1 2 max {0, log (sa/g log (1 (1- xx(x,xxxx a)))} such that g(x((11) cg. [V s ((x((x())) S Cg. and (x(*) Josish B(x("),r).
For any 17 > O mink, seste, (f(x(1)) f²) S (n+C'as)² if and only if > t1 + At(n, with At(n, == (ist?(x(!1),00)) Moreover, assuming Assumptions C.5 holds, we have f(x(()) - -f- 9(x(0) g" B(h(x(t)) - hg) aftn and so ming, (h(x(1)) -hg)) x+C'a8 if and only 12 t1 At(
2
Proof. Set = X(g, ((),)))))) r). Since (0) < x. there exist E (0, and E (0,c) such that 10 <
1--e r²x
149
(see Lemma C.11). Let Li- == ()) g < co, L2 == SuPxEB(x!0)) 2r) vec (V*g(x)) 11o0 < 0o.
:= vpLz and Ln == s( (upHeat(x) || < co. We choose any 3 E (0, 3max) with Bmax > 0 and Smax sup SuPHren(((0)) |||| S ())) and any step size a E (0, amax) such that
O'max min {
2(c-7)
1
(L1+ + Amax Ln)
L2p "(1+))
Let 8 (1 () x.x.xxx.. x E 10 (Lemma C.20) and 1 (1 7)aT E (0, From Theorem C.14, when
all V BLLn --, ; we have xkk C B(x("), r), g(x(()) S (1 5)*g and (+(((()) < (L)(x( <
2L(1 - 5)% (0) for all 0 < tt. Setting (1 - g I Eg gives
t: :: :==t(c>) 2 max (0. log(cy/g(0)) }
log(1 6)
0
if cg > g()) otherwise
(8)
loglen/9
log [-e)xx(xxx0 r)a 8(oxx(xx(00
log(e,/990) (2)
Sa 1
log 1 (1- xlg.x a)
From 7 (x) Ln 2 T r we have t I (Vg log(T). So for t(Eg) to be well define, we need
log(c,,/ /g (0)
( V "0) [|اا when eg < (0)) This is equivalent to e, > <g D3C with (ng((r)
lng( ()
lo
3Ln
log
E (0,1) since BL. < > V (1)
C log(1 6)
since
5 E (0.1) and 7 E (0,
(9)
log
LA
c
D = (0)
7 [x
21

---

Grokking Beyond the Euclidean Norm of Model Parameters
Let F(x) xx())= Vg(() + (30(()). From Lemma C.12, we have
(f(x(()) - dist. (x((x), ( + (12 - ti)a² naxi, [
a max F ((x(
min hitSt2
(10)
2(t² - aa
12-000
2 5
Since ||F(x()]13 O for all t 2 t when |[ < (x((()) (the second phase of training is driven by (x(()))), see Theorem C.16), we obtain from Theorem C.13 that there exists C' >0.
(n 3 - 22211 +
dist (x(()),))
(1(x())) - -x) S
min s
(LI)
2
aBn
And if , 0 no +0,
(M(x(1)
- Nig) 11g) -
17+
C'aB
dist
min Ststta
- 22+
(12)
2
a
Lemma C.11. For all a, b E R with 0 < a < b. (1) There exists c E (0,1) 1) such that a < (1 - (1-c)²b. (2) There exist
E (0,1) and y E (0,e) such that (
1-
b. (3) There exist (0,1) and y (0, - c) such that
([(ett))
Ity )
b.
1+7
Proof. Set = 8 E (0, 1) and T = vP E (0, 1). (1) Take any c E (0, - r) - p < (1 - c)² < 1. (2) Pick
7 1-s 25 E (0, | and E (0, ( ). (3) Pick 2(1+T)
I-r E (0, 1+7) and E (0, -r)
Lemma C.12. For all t and t2 with t S t2, we have
(12 t2 maxt, [I
min 6 SrSt
(5(x(0) - - 1
a max 1i
| [(x
(13)
2(t2 - t
22-000
and
sup (x(0= xx) G(x(9)= + (m(x(()) - -n3)
[s(x(()) -0r) + 3 (h(x(()) - ni)]
min 4 StSt
S min
xEB,
s
dist? dist?(x((1)) O,)
,g)) + (t2- - +1)a² maxe, S15te (x)
2(122 11)x
(14)
Moreover, if Assumptions C.5 holds, then
(s(x(0) (*) s dist (iis(((((t)) (x(1)),0)) + (12 11 )0² maxe, sest, [[ ((s S15t
3 min (M(x!0)- - -ma) S min 451512
(15)
2(t₂ 11
Proof. By the definition of the subgradient (()) of f at x((2), we have (x- - x(())) T F(xe(ta) S f(x)
f(x(ta)) -x=--(xxx - x)) T F(x((a)) < - (xx(x())) - f(x)) Vx. So, for all x, we have
[x(+++111 x112- = 11x((2) aF =x((x()))= - -x12
= 11x((2) x112 - 2a (x(ts) -x) T p(x(lal) + (x
$ (xx( -x12 - 2a (fxx(ta) - f(x)) + + (2
(16)
12
12
$ (|x(()) x13 2a
Σ (F(x()) - s(x)) + a² Σ ||F()!0)
t=h
1=
22

---

Grokking Beyond the Euclidean Norm of Model Parameters
This implies
f2
0 5 inf [x((ttt) xx11 S inf (|x(t) - -x1 2a sup Σ
(1(x(0)- ((x)) + + E (1
x
x
xEO t=1
1=
1/
12
= 2a sup
Σ (f(x()) - f(x)) s dist (x)t)) (x)(()))))))))) ,)) + 2
|| 1
x t=t1
(17)
1=
= 201(12 - 41 sup
min xEe, 151512
(s(x(()) - s(x)) S dist fist?(x(!)) (x(tt) 0)) + (12 01)0²²
max [|
115152
(s(x()) - inf s(x)) - dist (xst) ,y) + (tz - t t2)a² maxe, Sesn, [(x(()) 1
min SiSt
xEE,
2(1²- t,)a
This proves the first inequality. We also have, for all x, - -xx(tx) - xx) T H(x((ta) S -(h(x(lx)) - -n(x)) and -(x(()) x) G(x(!t))= 1 - -(((x(()) - g(x)). So
[xx()) x112 [|x(12)= - aF(x!(2)) - -x12
T F(x(ta)) + a2(11((((12))11
[|x(ta)= x 2a (x11)) -x)
[|x(ta) x13 - 20 (x((x) )x T G(x(tt)) - 2a,3 (x((x)) - -x x) H H(x((z)) + a ( (18)
S (1)) -x1 2a (0(x(t))) -s(x)) - -20B (i(x(t))) - -xx)) + a² (
12
< [1x(0)) -x13 2a
(o(x(()) -s(x)) -2ag3 Σ (n(x(0)- - -(x)) + a
Σ 1F
t=
t=
t=
This implies
12
12
2a sup (((x(0)) -100)-990) (000))-9(x)) )(xx) + + sup
Σ (b(x(0) - m(x)) S dist? (ist?(x(1)), ,0,) + a²
> [|
x t=b
xEB, t=t1
1=)]
(19)
[ - 9¹) + 3 (h(x(") ni)] S dist dxst?(x((11,, (x(()), + (t2 aa² max,, Siste [Fx(((1)) >
min 4 S
2(12 - t1)a
If , ng, y, then f(x) - 2 8(h(x)) - -hi) (Lemma C.6), and so
8 min S155
(M(x(1))- -h₂) S
min
((x)()))(*)
min Gsrst
(fxl(0) - inf f(x)
(20)
Srst
x
Theorem C.13, Lett > O. Define R = dist(x((1),)) and L == max, Si [[( Assume there exists a constant + n²
C 0, L CB2 Then, for any n > mini, (f(x(1)) - -f*) S CatCa if and only if l1 Assumptions C.5 holds, then mine, (h(x(1)) - hT)= S g+Cag if and only if 2 h +
Moreover, if
n²
2
Proof. Using maxe, StSt2 ||F S max S [x(((() )13 = I S C,32, we derive the following from Lemma C.12: R² + C(12 4) 222 2(t² - 1
(f(x(T)) - -r))= 5
(n Ca.3) 8
R2
113
R²
min Gstts
1/ - 4 2
(21)
2
2a - -12
2
aBn
and
min (M(x(0)) -n;) 5 3
R² + C(2 - 2²) 3²
<
17 + Cap
R²
<
7
-
12 4 >
R²
(22)
G5156
2(t² t1)a
2
28a -22
2
a:Bn
23

---

Grokking Beyond the Euclidean Norm of Model Parameters
Theorem C.14 (Time of Memorization of fixed-step gradient descent under CL condition and subdifferentiable regu- larization). Let g RP + oo) be a nonnegative C2 function and h. RP [0, oe) be a nonnegative subdifferentiable
function. Take any x(") E Re with g(x(0)) + 0. Assume that g is r-CL at x(") E RP for some T >> 0. Choose
2
6 EE <E(0,1) 1) and 7 E (0,e) such that 4g(x(00) < I+7 ,x(y, (0), r), which is possible since 4g(x(")) <, xr2, x0))r)
(Lemma C.11). Let L1 == B(x00)( [V < <00, L2 == supxe 2r) vec (V'((x)) 1l00 < xo and
Ln == B( (ax)) | 10o < cc. Define G(xx == Vg((x) and choose 3 S Bmax with Smax > and Ile- Smax (UD ||11|| S (x| [1)2 Choose any step size or > O such that
{(En ((1 + Bmax L
2(6-7)
1
} with L. = vDL2
a < min
(23)
Lep " (1 *(1+))
Set X == x(gxx(),, .r) and define 8 == min - e)xa}. Iteratively define x( = x(k)
a (G(x(k)) + ()) y (x E oh(x(()) for each ks 2 O. Let Tert== == 1 1-
(1 11
7)l E (0,1) and assume
[|]G(x(())]zz 24n for some k. > O. Then
BIn >|10(x01)12 2 7 and & sup H
|| S (1) VOłjsk
(24)
and
[GG 2 TBLn with T8Ln < BL.
(25)
As a consequence,
x(1, x(),...,>x*** (8) E rr)
g(x())) S (1 - (0) (0) and ()x s (lg(x(1) s 21(1 - -)'g(x(0) for all 0 S s k
[|x(() - xx00 s (1 7)² for all 0 S j < h
(0)) since x (0)) E B(x("),), r). So (1) E B(x("), (Lemma C.17), and hence ())) ||11|2 S Ln. Since g is C²2 r(g(x) is symmetric for all x E IRP and we thus have max ( (x)) S
Proof >I|G(x(0))122 2 flu > 3L, 2 (x)
/pmaxi (x2())) S vpl2 = L for all x E r) c B(x("), So g is L-mmooth on B B(x("), r) by
Lemma C.4, i.c.
[IG(x(1) - G(x - L1(x(1) - - 1l2
= =aL|G((x (0)) + (x1
(26)
Sal. (I ||2 011H(( (Triangle inequality)
$ G1+)) (0) 112 since (1| ll) S >1 ()))
So
|G(x0))|2 2 - (() (Triangle inequality)
2 ||G(x 12) - (1 + +1)aTT((((((( (0) (Equation (26))
T||G(x(0)) [|2 since T -= - (1 + ))L
(27)
3Ln
since +|G 2 LLr
-
7
We thus have >|[G(x())]l2 2 BL. > (x))) So x (2) E B(x .r) (Lemma C.17), and hence sup sup | 2| BLg - and (supheah((a)) || ||2 S > G((())) [l2 V S k. We have |] 2 TBL, with TBL,, < LL, So either [G(((++))]] 2 BL, (and we can proceed to the next iteration), or 3L, 2 [[(((++))) 2 TBLn, (we do not have control after k). The last part of the Theorem follows from Theorem 1.
Ln. And so on, we prove that ||G(x || 2
24

---

Grokking Beyond the Euclidean Norm of Model Parameters
The following Theorem is extension of the result obtained by Chatterjee (2022), where we add an additional bound on the gradient. We extend it below handle the case 8 * 0 in Theorem C.16.
Theorem C.15 (Convergence of fixed-step gradient descent under CL condition). Let g. RR + (0)) be a non- negative C² fimction. Take any x(") E Rp with (xx(")) # 0. Assume that g is r-CL at x(0) E RP for some r > O, ie. 4g(x(0)) < x2, x()) r). Choose e E (0, 1) such that 4g(x(()) < (1 - x(),xx)) which is possible since (x) < xxx(xx(00) B(x() 2r) || vec (V2)(x)) bl < ca. Choose any step size a > O such that a < min Live Lap } and iteratively
(Lemma C.11). Let L1 == su r)) [Vg(x)llo < oo and L2 ==
2c
define x (l+1))= (k) a (g(((()) for each k. 2 0. Then,
x (M) E B)x r) for all k, and as l-- + oo, x(() converges to point x* E B(x("),r) .r) where g(xx) = - and 12
Moreover, for each k 2 0, |bx(*) - -x113 s (1- 8) ()x S (1- (x) (1) and [9(x(())1 2Lg(x(()) s
21(1- -
with 5 == min (1, (1 - c) x(1, (0) .} and I. = vpI.2.
Proof. This is a special case of Theorem C.16. The proof sketch is the following. Since | - x x000 122 S Ei=} ||x(++1)
x0112 = 3 all vllVg(> x [12,, Lemma C.17 use the fact that, if for some 2 1, x(1,, ,..,,,,,'' (,,,...,x((t)) ((1-1) E B(x (0) .r) (this 19(x(0)
need prove Lemma C.17 by induction), then Σis a: (g(("))) S (1- 5)0/2 v x(g.x(0), (1--33) for all 0 و
1
(Lemma C.22); so that [|x(A) - -x(0) ||2 r since 4g(x(")) (1 - e)3p² x(g, (0) ,r), i. (k) E Bxx(()) r). Lemma C.22 on its turn uses al (g(x(() ( S g(x(0) (((x((+)) (Lemma C.21) and g(x(i+1)) 1 (1 - (1-5) gx(((()) (Lemma C.20). Both Lemmas C.21 and C.20 use |2] S ea:l (x( (Lemma 1.) with R = g(x(+11) - g((())) + (
Theorem C.16 (Convergence of fixed-step gradient descent under CL condition and a subdifferentiable regularization). Let g: RP + [0, oo) be a nonnegative C² function and h : RP + [0, oo) be a nonnegative subdifferentiable function. Take any
x()) E Rp with g()( + O. Assume that g is r r-CL. at x()) E RP for some r > O, ie.
Ag(x(0)) < r² x(1
(28)
Choose c E (0, 1) and:
(1) E (0, c) such that
4g(x(0)) < <11-g (1+7)
1² x2((xx(
(29)
(2) E (0, 1 - c) such that
4g(x(00) < 1 (1-(et22) r² xxx))
(30)
The two choices are possible since Equation (28) holds (Lemma C.11), Let
In == =UUXxEB(x00, r) [VVx()) < co be uniform upper bound on the magnitudes of the first-order derivatives of g in ())
Le == (IPxeB(x()) 2r) || vec (x)(x)) Ilo < oo be a uniform upper bound on the magnitudes of the second-order derivatives of g in ()x(0),2r)
and Ln suPxe .)) SupHEоM()) || < co be a uniform upper bound on the magnitudes of the first-order subderivatives of h in ())
Define G(x) = Vx(()) and choose B S 8max with Bunx > 0 and Biax sup (eB(H
||1]|2 S (1] Choose any
step size a E (0, mmax) with
min
2(5-2)
(1)
(Ta Tnlpp' Lgp
s
a'axx
(31)
min [ ma
Lzp
2c
(2)
25

---

Grokking Beyond the Euclidean Norm of Model Parameters
Sel x= = x ,)) and define
min (1 e)xa} min (1 (E x))
(1)
ô =
(32)
(2)
Iteratively define x(++1) x(k)) (G(x + (x(H(x((1) A H(x(1)) @h(x(1)) for each &s 2 0. For all &s such that
(supheσn(Xá)) |H S >][G(x())|lz 0A S j < the following hold:
(8) E x(),,))))
rr)
s(x(i)) < (1 - for all م
[Vg(((()) 113 s 2Lg(x(0)) 21(1 6)1g(x(0) for all S S k. with L vDL2.
||x(*) -101 s (1 6/ for all 0 0
The sequence (x(0) losj<k Cauchy.
Moreover, if (uuphean()))) ||]| s 7]|G(x6))]l2
Vj
7j20 0 B = 0), then {x(*) converges limit x*
B(x("), where g(x*) = and [Vg((*)]l = O, and we have [|x(M] x"1 2 s r(1 - (1-8)k/2 for all k 2
(1) 122 A 0 s < k. We prove in Lemma C.17 that Proof. Let k 2 1 such that 8 {u || S <>I|G((x0)
x(),,.. ,,,,, (k)
for all 0 < k. Since g is C2, (x)(x) is symmetric for all RP, So have for all B(x (0) .r) C (() 2r), Amax (x)) max (V2g(x)) < is vpL2. So is L-smooth ((x(") r) by Lemma C.4. This implies [ - 2Lg(x) for
B(x((),,)) r). We also prove in Lemma C.20 that g(x(0) S (1-0)
mamaxi [V((())
all (xx (0), rr) by the same lemma. So V s 2lg(x(0) s 21(1 8) (xx for all 0 S j $ k.
For all s k- we have
(xx 11
[x(())) -x1'112
a: al|((xx!') + (3((x
[=j
(33)
(1 (1-5)//2 V /41+7)9
if S 2(c T22
$
x(1-e]²
(Lemma C.22)
(1 - -0)0/2
if S 2e/Lap L2p
((1-()+>))*
r(1 0)//2 (Equations (29) and (30))
with
5 6:= {
(1 =)xa
if < 2 ((e-7)/L2p E (0, (Lemma C.20)
(34)
(1 x)))
a 5 22 L2p
For all > 0, if we pick any
i =
1 2og((s/r) ToR(I-) 0)
if - 1 5r otherwise
K(s) max {1,
2og(s//) log(1-))
otherwise
(35)
then we have || (k) 11 < & for all j. > K(e). So the sequence (x(()) ( Josseck is Cauchy.
sup || ||]| < >x Vj 2 0 (e.g 3 = 0), then (x(x) Ja20 Cauchy, and as a Cauchy sequence in
the closed subset (x) r) of UR it converges to limit x* ((x((),)) r) that also satisfy ||x* - x ||2 S r(1 8)5/2 for
all j 2 1. Finally, taking k + oc in Lemma C.20, we get ( ()) S (1 - 6) gxx())) for all j 2 0. This achieves the proof of the Theorem.
Lemma C.17 (Iterates stay inside the ball). For all lo 2 0 such that B (uup))o)((())) || S >1|G((x))) 111211111 0A S < we have (1) ..... (kk) E (1x r).
26

---

Grokking Beyond the Euclidean Norm of Model Parameters
Proof. We prove that ( E (xx("),, for all & 2 0 such that (ea(2)))((6))) || S +||G(x0))|z 04 S j by
induction on k. The base case k = 0 is obvious since (uu | |12 S Pmax (o) ||A|| S <>1|((x)))|l |||
Let k 2 1. Assume x(6) B(x(0),,) .r) and (su || 1 >l(G(x())1 for all s j I 1. Then
x(*) E B(x(0), r) since
|xx(0) x 112 S
alIG(x + (1(x
((1+))
a S 2(E - L2p
x(1-e]2
(Lemma C.22 with 0)
4(1+7) x(1 (ctt(1+)))
if < 2c/ Lyp
(36)
a S 2(E (((-~)) (cc-n)/L23 L2p since 4g(x(0))
,2x (Equation (29))
1+7
a S 2e/ L2p since 4g(x(")) <
(1(ct2)
r²x (Equation (30))
Lemma C.18 (Second-order Taylor formula with integral remainder). Let R + R be a C2 function. For every base point xe IP and step he RP, g(x + h) g(x) + Vg(x) h + fo - f)h Vig(x V + th)hdt.
Lemma C.19 (Second-order remainder control). For each j 2 0 set R2 = g(x((+2)) - -g((()) + al|G(x(1))13 Suppose
that x E [(x("), .r) and Ssu (reatea)() ||A|| I >|| for all 0 S k- 1, for some k 2 1. Then for all
0 S j < k - 1,
|Ryl 5 { Eall callVo(((0) (e (e+))al (
ifa S 2(- - L2P
(37)
ifa S 2e/L2p
R, and let x(t) == x() to (G(xl)) + (HH(((1) with t E 10. Apply the second-order Taylor formula with integral remainder (Lemma C.18):
Proof. Let Sy g(x(++)))
- g(x(0)) + +GGx(()) + +(H(((0)
Fix < k- -
((x(++1)) = x () (i( + 31((((1))
+ a² L'' Vo(x6)) (x(t)) V9(x(0) (x(x))))) -tydt
(38)
= -)(x(0) aG(x()) (G(x())
+ (x(()))) (G((())
Hence
5 02= //' (v
t)dt
(39)
Along the segment x(I) with & [0, 1] we have x(0) x(0 x(1 x(+1) all points x(I) lie in the line segment joining x(i and andx( x()+1), Because [x0++1ix((1)) S v - 1000 (
+ s
r, that segment stays inside ((x("), 2r); thus every entry of (t)(x(t)) is bounded by L2. Therefore, for every t,
Vo(x())V<o(((00)
(40)
L2
[ |(((9)111 L2 Vo(( < Lepl Vg(x(11)112
So
|8,1 18 22 //'
/ (0gx)ba ( (2((((())V 11-6 t/dt dt
2 (22220199(x(1) /.' (1 l)d
(41)
8²
(E (e- (a- c
if < 2(c ((e->)/Lo L2P
= - Vap 11
if < 2//L2p
27

---

Grokking Beyond the Euclidean Norm of Model Parameters
This implies
|ayl = |S; - ) (((())))))(())))) s | 1ee
< +5]]+ay||G(x))) 11) ~ (e +
{
i S 2(E - L\
(42)
+
a S Be/Lay
Lemma C.20 (Geometric decay of the objective). Let
6 :: {
(1 - xxxa
if S 2(c - ~7)/
(43)
(1 - xa
ifa S 2e/ L2p
We have 5 S and ifx()) E B(x(0), r) and B ( | S 7 [|| for all 0 sSisk k for some k 2 1; then
9(x(0)) s (1 ()) v0 م
(44)
Proof. By the definition of x= == X(g. x x(g.x(1)
(x)))))),))))))))r))) .r), we have xg(x) 1A Vg(x) for all x E B(x((),) .r). Since x(6) E B(x(
for all 0 S j -
1, VV x (1) 112 2 xx for all $ j In k 1. So, for all 0 I j S 1,
9(x 9(((0)= a VV 11 Ry by the definition of R2 9(x - a VV
1 + [allVg((())]
if a S 2(e - b2p
(Lemma C.19)
9(x61) - a V
1 + (E+ ) V
a S 2E L2p
{ 9(x(1) (1 -))l V
if S 2(c 7)/ L2p
(45)
9(x6)) (1 - (e) al [g
113 if - 22/Lyp
9(x(1) (1 -e)xxx(((())
if 2(e 7)/ T2p if L2p
9(x(1)
(1 (c + 7)) (xx(x(ii)
= (1 - 6) 9(x(0)
1- Moreover, taking j 0 gives ()x S (1- g1))) - 82 ((x(1)))g((()))))) ((x)))/((())===
Lemma C.21 (One-step descent bound). Suppose that x(6) E B(x(0),,) r) and 3 SSuPHEon(KO) | S >)| for all 0 < j < k - L, for some k Z 1. for all 0 S S 1
g(x())) g(x(+2)) 2 { (1 (1-e)allVg((()))
if S 2(e - 1)//2p
6 [
if < 2e L22
= x
(46)
Proof. For all 0 sjsk-1,
(xx)))-(xx(1)) - [ (1) )12 R2 by the definition of Ri
a (())) || eal| )
a < 2(e-7)/La) if < 2c/Lop
(47)
(Lemma C.19)
a [g(x(0) (E + > V
Lemma C.22 (Bounding cumulative step lengths). Suppose that x66 E B(x("),)) and Bsupheanceon ||]l S
>|]G((())) la for all 0 < < k for some k 1. Then for all 0 I x < k -
{
(1 - 83//2 (1-6)//2
4(1+7)2gxx00) x(1-t)
if $ 2(c - /)/L22
() + (HH(
(48)
(1 - --//1
4(1+7)Pg(x(00
if S 22 L2P
x(1 ()+t)))
28

---

Grokking Beyond the Euclidean Norm of Model Parameters
Proof. We have
alIG(x!0= + +31((x))l: )l2
IMI EMMMMMMM
a (IG(x()))= + +111(())))
a(1 + 1)] G(x") since 3||H(x s <>11)(x()|12
a2(1 [11++0 +11+++)110(x000111
a² + 3)x (u(x(0)) ((x(+11)) (Lemma C.21)
8
a +
Vo(x()) - Volx(+1)) ( (Votxi)) + Volx(ii1)
8
1 (Cauchy-Schwarc
Q + P
Σ (Votcto) - V
Votxxoo) + volx
8
Vo(x(0) + v / )
Q2(1 + 7
(Volxc)) - (o(
(Telescoping the first sum)
8
(49)
a2 + x)x
Ve(x6n) Σ
1-1
2 (
since g(x S (x) 14 < k 1 1(Lemma C.21)
5
20
x)3x
g(x(0)) (1 87 Σ
V(1 5)
since g(x(0) S(1 - (0)) (s))) (C.20))
8
21 + (xx(((()
(1 (1-8)0/4
8///4/
Σ (1 8)112
6
21 + (x((((())
K-jz1
-11-1)8 -118))144
(1 (1- 80//2
Σ (1 (1-88) 6)1/2
5
1=0
20 + 7) ()
k k-j-1
S 5(1 - 8//4+3/4
Σ (1- -/-)
since vI 1-1-6/2
1=0
2a (1 + 7) (x)
(1 - 6)5/2/
(1 - (1 11-1111//// 6/2)
5
2a (1 + xx
< (1 - 6)812
(1 1 -8/112
1(1+)) (1x(0)
for = (1 - e)xa
x(1 13
11-8)112
(1+>)'9
for 8 (1 (e + 7)) xa
x(1=((+))))
29

---

Grokking Beyond the Euclidean Norm of Model Parameters
C.2. Main Theorem under Polyak=E.ojasevvizz Inequality
In this section, we prove the geometric convergence for g under the classical Polyak-Lojasiewice (PL) inequality, Definition C.23 (Generalized Polyak-Lojasiewice inequality). Let yp RP + IR be sub-differentiable function such that arg xeR (xx + ( and "n := minxeRo (x(x) > -00. For j > 0, y is said to satisfy the p--PL inequality if and only if
(x() p- - 2f1 aEdp(x)
inf
| 2
Vx e R
(50)
The PL inequality is typically stated for differentiable functions. This version allows yp to be subdifferentiable, which is a generalization, and thus is not always equivalent to properties known for the smooth case. The CL inequality can be seen as a strengthening of the PL inequality and is related to the classical Gurdyka-Lojasiens
inequality (Chatterjee, 2022).
Theorem C.24. Let 8 R2 + [0, oo) be a subdifferentiable function such that arg min ((x) + 0. Assume x ==
millxeR* (xx = O. i is for some μ > O, then for all x E IR, yp is r-CL at x whatever y² 2 2x((x)/u.
Theorem C.25 (Convergence of fixed-step gradient descent under PL condition). Let g RP + 00 oc) be a nonnegative
C², L-smooth and μ-PL function such that Gg == arg min g(x) * 0. and h : R. - [0. oc) be nonnegative
subdifferentiable function. Let g² == minxee= g(x). Take any x(0) E RP, choose any step size E (0, T.) and any 3 2 O.
Iteratively define (1++) = x(t) a (G(x("))+ + (x (x( for each 2 O, with G(.) = Vg(()) and (-) E @h(-). For any
y E [0. 1) such that k(7) = (2- - -aL)(1- - 7) - aLn? E (0, (0,2/pa), and for any t 2 0. we have:
B1|H(2!9)12 S >/|C((x())]2 = g(x((+1)) - g" s (1 - a x(>)) (o(x(0)) - g*)
(51)
(1 S {( Ak <t = ( - 9-| s $11 μ a n x()) k (9(x") -
90)) Wk St t (52)
As a consequence, if Equation (52) holds for some y E 10 with k(7) E 00,2/pal and for all t 2 0 (e.g. 3 = 0), then:
((x( - g" S (1 - p. a k(7))* & (9(x(()) -r") wk 20
= (1 - a. (2-ul)) (9(x(0)) - wk 2 0, for 0
(53)
= 1 -(4)
(1(x((0)) -
( 5r) wk 2 0, for = 0 and a 1/L
Proof. Fix t 2 O, and assume (1(((((0)|l I|12 I >I|G(x(0))2 l2. Since g has an L-Lipschitz continuous gradient, we have
((x(++1) s {x(x(()) + (x)) + 5 -1x(((+)) (Lemma C.4)
La² |[G(x(1)) + (x
=9(x(0) - a(G(x()) + (xx T G(x(1)
2
= =9(x10)- a (2- - aL) |G( _ " (2 - (L-L))31 aL)3H((x()) (1) ((x(()) + LLL (x >
a
2
2
2
a²L
s =x(x
2 (2- æl. |G 1
a
2 HT
((x(0)))
|3 (x )0
2
2
S 50(
a (2- - al) 1)
a 2
1|2 +
a²L
13H (x(x))1
(54)
2
2
2
S 90
a (2- - aL) (G(x )
a (2 aL) > [ 1 a
aal ( "(x
2
2
2
= () - ak
2 [( with x = (2- aL) - (2 - aL) 7 - aLd 2 20 0
S 9(xx0) apx ()(x(0)) - since x 2 0 and و is μ-PL,, i.e. |[ (0) 2 2n(g(x(0)) - )9")
(1 apx)) (9(x(1) - )rr +9"
As a consequence, if we have 3 (x(())l )|2 S > for all t, then ()x - 9" S 511 (1-) (9(x)) - -g".
(g(x!?)) - -g"1 < 11 - aqulll((x ) 9-1 < |1 - apuas | (g(x") - g"). ..., (g(x(()) - ge1 < |1 - aux]|l((x*1) 11) - 66 S
- -aauu]] (g(x") - g)). Note that de (1 na (2 - aT.)) = + 1/L. L.
30

---

Grokking Beyond the Euclidean Norm of Model Parameters
One recovers the standard convergence result under PL inequality (Karimi et al., 2020) from this Theorem with | 0 (Equation (53)). Note that the smoothness of g along the trajectory is only a consequence of the result under CL., whereas under PL, it is assumed in advance. More importantly, the CL is only required at initialization, whereas PL is required on the entire domain of
1 = aL E 20 and n² = (2-7) . We want (a; 7) E (0,2/L) x (0, 1) such that 0 & x(7) - p/ We have S k(7) =
= 0. T v ((2-))2++3)-122-))
= +y²
(2 (2- -2
< 1. Since a. is downward-
27
opening parabola in y, its maximum value in the interval [0, occurs at 0 (since the vertex is at a negative y value),
Therefore, we only need to satisfy the inequality x(7) v 2/ at y 0, that is 2 aL I 2 - pLa² 2na + 2 < 0.
poo
The solutions for a depend on the discriminant A = (-2p)² (2) 4μ(p 2L) of this quadratic.
If μu 2T, A < 0, and the quadratic function μLa² 2pa + 2 is always positive. So the set of solutions is empty.
If 2L, A O, and the quadratic inequality for a is satisfied only at its root
2 1μL
+. For this single value of the
((2-))2+3)-(e-1) v5-1 So, the solution set is a line segment, { {G),) 1 e
[0. v5-1
condition on 7 is: 0 S y S
2(1)
2
2
If J 2L, A 0, and the quadratic inequality 2La² 2pa +2 S 0 is satisfied for between the two real roots
E
N 122 2al
a+ n p
C (0,2/L). For any a chosen from this interval, the corresponding values for y are
ah
pL
given by 0 STS
(2+
-2+3a (2-
2oL
C.3. Preservation of the Chatterjeu-LLooasieviie Inequality under Perturbation
In Section C.1, we study how g behaves when a regularization h is added to the optimization objective f= g Bhh. We also study the behavior of f and the regularization term h when g and its gradient are already small. We establish below that, under some assumptions on h and g. if g is r-CL at x E RP, then for all (0, 1). there exists Bmax maxx x, r) > 0 such that for all B < fuax: the function f = g + ,hh is also er-CI. at x. This means that Theorem 1.15 also applies to f*. In other words, we also have (undefined terms are those of Theorem C.15):
(A) E B(x!"),,) .5) for all k and as k + 0o, (k) converges to a point x E B(x("), ,r)) where f(x*) = f
For each & 2 0, [xx( - 1 < (1 - 2) and f(x(1)) - - S (1 - 5) (f(x(0)) f*) with 8 == min <) x(f, (x)) er)) a}.
We start by introducing a property less restrictive than convexity, which we call the Zero-Stationary Property (ZSP). Definition C.26 StttonStatinnary Property). A subdiffcrentiable function yo Rp + R is said to satisfy the zero-stationary property (ZSP) or to be ZSP on U C R2 if and only if for all x E U, O E 0p(x)) = ((x)= = 0.
For a ZSP function, if a point is stationary (in the sense that O is a subgradient). its value must be zero In the context of non-negative functions, this implies that only points with the optimal value (zero) can be stationary. This is linked to Fermat's 's or interior extremum theorem (which identifies stationary points via derivatives), but adds a zero constraint on the function's 'ss value. The Lp norm (p > 0) and nuclear norm &, are ZSP on RD. Convex functions are not ZSP in general. Take for example p(x) = 1 or =x() = (x- 1)² 1. But if yp is convex and ypt == inf inf sp(x) is in the image Im(Ø) of yo then y is ZSP (Lemma C.27). The ZSP property does not imply convexity. The ZSP is local because it only constrains the function's 's value at points where the subgradient contains zero. It does not constrain the function's 's curvature or behavior away from these stationary points. Therefore, for a ZSP function to be convex, it must satisfy one of the standard conditions that define or characterize convexity (first or second-order condition, Jensen's inequality, convexity of the Epigraph, etc).
Lemma C.27. Let : RP + R be subdifferentiable functioo such that -0 < xp := infxeR" xx() E Im((p). If up is convex, then 4p- x0" is ZSP The converse is false.
Proof. Let to be convex and assume -0 < 4- = infxeR p(x) E Im({p). Define ((x) := p(x) - "0" Let E IP such that O E ap(x)= = Öx(x). By the definition of convexity, we know that ( 2 (x() + (y - x) for all y and all z E @p(x)). Since O E 0p(x), we choose z = 0, and get p 2 x(x) for all y. This means that x is a global minimum of
sp. Since f0* E Im(pp,, we conclude xxx) = pp" i.e. ((x) = 0. This proves that و is ZSP.
31

---

Grokking Beyond the Euclidean Norm of Model Parameters
Take (2())
x² 1+23 for R, with 9" O. Since ('(x) (1+z2ya we have dp(()) = 0 = f(z)
So
9 4- is ZSP. But y is neither globally convex nor globally concave on I because ("(x)= 2(1 3(1-3x²)
changes sign. LI
(1+02ys
Theorem C.28 (Chatterjee
stability under 7SP Perturbation). Let g RP + [0, oo) be a differentiable function, and let h: RP + [0, oo) be a subdifferentiable function. For all RP and > if g satisfies the r-CL inequality at x, ie xg(x)= < 2 x(g.x.r), then:
For all & E (0, 1), letting r' == er, the function = 9 + Bh satisfies the r'-CI. inequality at x with respect to g. i.e., 4x)x) < 12 x ,g) provided that B S 3max(e, min Vgxioxx.))
Moreover, if sly) == fly) O, then f 9 + Bh satisfies the r'-LL inequality at x provided that 3 satisfies this constraint, i.e., 4f(x) <2 x().x.r")
-/-BB 2C
41.
h is ZSP, then for all & E (0,1) such that 4h(x) > r x(,, (,,h,,,"' x,,') with r' := Er, the function f := + -Bh satisfies the er-CL inequality at x, i.e., 4x(x) < < x(S,x,er) provided that B I Prnaxe (e, x, with
V9.x(9.x.7) x(g.x,r)
VA --- B
4g(x) "112 x(),x.1") 4h(x)
= min
(55)
4L
2C
If is ZSP, then for all E E (0, 1) such that O & g (B(x, ER)). the function f == g + BBh satisfies the Er- CL inequality at x, ie., 4f(x) 62p2. x(1, x,er), provided that B S Brnnxee (E, (ex,,r) x,r) r) with
0 x(g, x)
VA
V Tx x, er)
Amnx x, min
(56)
4L
20
2 2lx(h,xx,1) er)
where
r/ =r:
g :== inf
g(y);
:= sup yEB(x,r"
gly),
h: sup hly)
B(x,r))) s(y)
yen(x,r))
L = sup
sup
| E (0, 00),
A = (x,(k,xrr)) + x,
yx (y)
(57)
1 D = z r² x(0,xxr)
A = 4x) g.
:: = 4g(x) -o+ 4h(x)
4:=4h(x) T
A = B² 4C D) 16(g(x) h(x)g)² + 8h(x)h y2 x(g,x, r) -20
Proof. Fix any E (0, 1) and define r' == We have B(x,r') C B(x,r) since r' <r. Let E B(x,r') with f(y) # O.
Write ff(y) = Vgly) + Bzy for a certain Zy Öh(y))
1. If g(y) * 0*, then since y B(x,r) r) and g(y) * we have Vg 112 g(y) x(g, x,r from the CL. constant for g. We also have || Izyyl 2 S L since L B(p. sup
|z 2- Using Cauchy-Schwarz,
(Vg(y). zy) 2 [ | l12 2 - -
(58)
So
[V = V1
+ + 12
2 + :2y)
> ||Vg(y)] v
[[Vov))ll -3LL)
(59)
2 Vglr) x(g.x,r) (V)()) x(g,,,r) -28L)
(Volx) x(e.x,)) - BL) 132
'0()) #0 A(y) * = f(y) * 0, hut the converse is false.
32

---

Grokking Beyond the Euclidean Norm of Model Parameters
Since 9 x(g,x,r) 2 43L by the choice of 3, we have v v x(g, r) 2 48L. So v Vg(r) x(g,x,)) 28L 2
28L 0 and vg(y) ((g.x.)) 23L 2 ( x(g, x, r). This gives [ 2 (g(ly) x(g,x,r).
Since f(y) = g(y) + Sh(y) - 3h, we get
[V l2
f(y) 2 x)
Ml 2 xx(g.x,r) ((3) with
gly)
g(y)
g
((3) =
inf (x f1))0, g(v)#0
2
(60)
fly)
9 + 3h
Plugging x(f.x,))) 2 x2x().x,) ((3) into the CL inequality:
4f(x) 4g(x) 43h(x) Er2 x)..
9
S 12 x(fxx,t)
(61)
g+ 3h
Multiply both sides by 3h, we obtain the inequality (4g(x) + 4x) (g + ph) and rearrange, we get C,3² + B3 + (A D) < where A, B, C,C, are defined above. The discriminant of this quadratic is:
< frr² x(g, g. Expand
A = B² 40(A - D)
= 16(g(x)h + h(x)g)² 4. (4h(x)T) (19(3)0 1
2" 2"0² x(a.x,)) g)
= h (2g(x)((x)) + 1G -4g0()66g)(((xxx)) S Sh(x)hr x(g,
(62)
= =16g(x)*h' - -22g (2) (x)
+ 161 8h(x)Tr x(g,
16(g(x)h h(x)g)²++++))))))) 8h(x)rr x,) . 20
Therefore, the inequality admits a real solution, and the optimal range of 3 is given by B < BBYA 20
So f satisfies the
r'CL inequality at x with respect to g. i.e., 4f(x) < 12 x(f, x g). provided that / < min Moreover, if gly) = 0 - [(y) = 0, then ʃ = g + 3h satisfies the rCLL inequality at x provided that satisfies this constraint, i.e., 4f(x) < <2 r')
Vgx(s.x.1)
VK-B 20
4L
2. If g(y)- = 0, then f(y) (h(ly) # 0 and zg # (since 0 E Oh(y) = h(y) =0). So
2
|,, hly)
z
fly)
8x(h, x')). Plugging x(f,x, r') 2 6x, x') into the CL inequality:
4x 4g(x) (x) < r x(h,x, S -(fxxx x(f.x,r)
4g(x)= < (1/2 x(h, x, x,r') - 4h(x)) 8 with 12 x2,, x,, -4h(x) >0
(63)
22 x(,,x,r') - 4h(x) {
4x
< B if 41 < x22((1, x(,,x,r')
> 3 if 4h(xx) > 122 x('x, x,,"
The equation 4x(x) r x2((,,,,,xx_ x, r') means h is r'-CL at x. There exists x such that the lp norm (p > 0) and nuclear norm l. satisfy this condition.
Assume 4h(x) < x2, r"). Since 3 < min possible if and only if
(axxlaxxr) 4L
YA-B 20
(Equation (56)), the choice 3 > "=.x((,xx7") 4h(x)
4g(x)
4x(x)
V@ x(g.x,r)
VA
B
< min
(64)
r') 4h(x)
4L
2C
Ensuring that this constraint is satisfied can be complicated. To avoid this, we will make sure that only the first case above is considered. We have
| fy)())
x(axx.)) J+Ah 8x(h,x,
if fyl) + 0 otherwise
(65)
f(y)
33

---

Grokking Beyond the Euclidean Norm of Model Parameters
So if we impose
1 x(axx.1/g 2 7+8h
2 8x(h, x,, we have
/ fly)
2
x(axx.)) and thus 3 <
-B+vA 20
The constraint is
3+8h
20
equivalent to
1
xx((,,x,"' 68
B 2 x, x,) 0
-VA' - _x(,x,x,r))
S 5 VAL - xx(
(66)
with A' = ( +2xx(sxxxr))
+
2,x(1,x,rr)
2hx(h,
C.4. Proof of Theorem 3.1
Let ya) = Xa. We have y" = Xa" + E, and want to minimize f(a) = -9(a) + +.8h(a) using gradient descent with learning
rate a > 0, where h(a)= = ||a|| and
1 X + Ell
((a) ly(a) y"112 2
x Xa (X Xa" 1aa+xx +x a
(67)
We write F(a)= == G(a) + (a) with G(a) == Vag(a) X (X x( and H(a) E O|al|| any subgradient of ||a||1, that is (a); sign(a;) for a, o O, and any value in |-1, for a, = 0 Suppose we start at some a!)), The subgradient update rule is
a (1+1) at) oF(a!") (In a x)aa (0) +a (X Xa* + + E) (aH(a()))
VA
(68)
To explain grokking in such a setting, we will look at the landscape of the iterate a!" (1) Let X UV V under the SVD decomposition, with Σ = diag(Tr) 0r+1 = = 0. We assume by default the SVD to be compact, i.c., U RNxr and V E Raxr have orthonormal columns, but we will make precision when we want full, i.e., they also orthonormal rows, with that time U e RNNN and E R Using E = I =Z,, the dynamics rewrites
where r = rank(X) and max 01 2 ... Tk 2 Tktl 2 mmin or >
a (1+1) VÊVT (1)) +a Va(VVVV'" a" VVAUTE (a3
We assume the step size satisfies 0 < a <
In fact, for the dynamical system to converge, we need (In, - a = 1- ao (In - aX T x) C (-1, that is 0 < aσ < Vk E [n]. For all p > O, let define Pp == |/ - axT p- so that /2 = ||, <>> 2-+2 = maxkelr] | - aorl E (0, 1]. We will show that for B small enough,
2
dmax
the update first moves near the least square solution of the problem, å x)) x y V + Σー UTE) with ( = -/5 (IN - TUT) x 1 11111. Later in training, H(a) dominates the update, leading to ||a ~ ||a*||1.
Theorem C.29. Assume the learning rate, the regularization coefficent and the noise satisfy 0 < a < amax
d
0 <<> 3 Amnn (X x) and |X "Ell S VCa3, > O. Let P2 == Fmax (Un. ax xx). There exist t < 0o and a constant
Vn
C" > 0 such that:
|1a(1) - ál12 2/23 and g(a ()) S y()) +
2na2 max (X)
1I 2 4, 2g(a) = =5 UU ) < ||
P2
(1 - )2)
Vn > 0,
min
(f(a(()) ( -(aa)) <
(n + C'a 8
t2 2 1 + At(n, (1), At(1,t1) == ||a (6) a
(69)
SSSS
2
a,31
+ (C+ (C+C")a5
Vn> > 0,
min
(Iaa0))] - ||a*1l])
<
1/ 2 4 +
SSSS
Proof. First, we observe that if 9 is too high, the subgradient term H (a) dominates early, and there is no convergence, i.c.. no memorization nor generalization. In fact, if 3 > Amax then the -tter dominates the updates, causing the sequence (*
to exhibit oscillatory behavior without convergence minimizer of f(a) g(a) + B|aaa||2 (Lemma C.33).
FFor the experiments, we used (a) = (ag.
34

---

Grokking Beyond the Euclidean Norm of Model Parameters
Let first evaluate g(ā) and ||a a² 12
We have
à= x
x
x y = v -
(70)
So
Xá = = Y vV a V VE VE - - - (lN UU
(71)
and
Xa y" 12 E (IN (U")(Ux UU J = E (UN UU " 3 - |
(72)
i.e. g(ā) 211 = - _ TUT & - E12
This implies, assuming E(g] = 0 and Cov(E) = of
2F (G) Fegl [á y"112 = Eell(Ix yy' = E { (IN - UU )
= tr ((UN - UU Cov(E)) + (EE)" (UN- - UU (EE)
(73)
= of tr (IN - UUT) (N -r)
Since
aa + â- = (xx) x -a' = v (v + + {) -a = (L, - -V
VE-
(74)
have
|á a a (I,, - VVT= a' x
(75)
Equations (72) and (75) show that the least square memorizes but does not generalize (for N < n). particular, a² has a nonzero component orthogonal to the column space Col(V) of V, then à cannot perfectly generalize.
From now on, fixed
such that Pp < 1, e.g. p 2. Recall G(a) Xy y") (x
(X Xa' + +) Starting from the update rule a (t+1) = alt) a (G(a(1)) (H)((())) we have (++1))
E).
(al0 - a) - a (G(a(1) + (a()))) Since G(â) = 0 and G is linear, G(al!) = -x X(a(i) - -)) Substituting this back,
(++1) = (² a) a (C(aa + ( (In - axx (a a) - -a)H(a!")
(76)
Taking the norm; applying triangle inequality and using [| S 1/8 give
(++1) - -1lp S prlla! - àll, + aBn 2/8
(77)
Repeatedly applying the recurrence,
(aa()) a|lp s o/(la(1) 11 al1p + aBn 1/p (1+ + Pp + ... + 1)
= p1 a|1p + a:n 1/p - P2
for Pp #
1 Pp
(78)
a/8n 1/2
S 05 allp + 1 Pp
for Pp <
Define
In ( (1 (1))))aa()))p)
All
0//01/7
t =
(79)
In(pp)
The definition of 4 ensures that for 2 4,
0' ('I(a()) - all, a:3' 1/1 1 1 Pt
(80)
Pp
Let u E a|a|[. Then [u:] S 1 for all i E [n]. |u||p = (E (1pp)) 1/r S 1/2
35

---

Grokking Beyond the Euclidean Norm of Model Parameters Thus, using Equation (78), have for all L 2 11.
|a(0) å|1p S 22,Bn 1/p 1 Pgz
(81)
1 Pp
Using this, we derive
Xa (1) y" = |X(a - â) + (Xá - -y")ll
< 1X|p-=1laa م + |Xá y'11p 11p
(82)
- 22 1/1 1- p/
-
X || ll-p + || y=11 11
1 Py
and
(G(a"))le |X ((a()) - -a)llp
< | T X|Ip-pp11 -ållp
1 Pg
< 20 //8
/
X X
X11p->p Pp
(83)
2a/3n
X X 11
1- Pp
2 2/73Bo
[X X) for p 2 since 1 P2 x
X)
and
1F(aa))2==1I((a8 = ))) (
||GG( (1) ]12 H((H(
nT.
(84)
X X 22-22 + 3/
P2
n (20max (x xx) + +1) 3
Also, since X(a(() a) E Col(X) = Col(U) and Xå y" = (IN - -U )g E CO)(U) we have for 2 f1.
s(a(t)) - g(h)= = IXa(t) -y113 2
Xa - y
2
I((a9) - a) + (Xá - -y) 11
Xa y"112
=
2
2
(85)
1
IX(a a)]l + 1 | y011
1
Xå y"
2
2
1
I((a)) - a)
2
So
2na² 2n 0
(X)
((a(1)) - (a) < 2
X |X|3
- -a12= <
(1-p2)2 O 2n max (X (X)
(86)
All this shows that after t,, the iterate (t), its error ()) and the associated gradient G(a behave respectively like å, g(ā), and G(ā) = up to an error of order (3)) Note that we assume 0 < Bn/n < F (X XX. So, the gradient of g can be made much smaller than the subgradient term after i by choosing 3 sufficiently small. After time 4, the contribution of the gradient G' to the update of a. is dominated by the e fj-regularization term up to an error of order (t) a aßH (a: a;) for all t t1. By Lemma C.31,
( Specifically for i. > I,, the update rule approximates a (t+1) 2 a (9) a (a()). If ((1)) > |a then
(aa (a, a²) (Lemma C.32), and a
(t)
(f+1)
ac a
(!)
this lead to |a (1) a!l < a for (and only for) t 2 4 +
ln
This suggest that we may have [|a(t) -aala=l O(a3)
aB
36

---

Grokking Beyond the Euclidean Norm of Model Parameters
|a -
for (and only for) t 2 1 +
In fact, considering the canonical subgradient (a) = sign(a), we have
a3
a a² = a()) a" -aG(a(t) a sign(a!)) Cap (t) a aaG(a(t)) with
2-7 2 2
i 7 if イ if -
Cg(2)= =
E
E [= - 1,,,,++!
k
i v m - 7) So the goal of (a) is to reduce
coordinate shrink operator, since [c,(x) - E [0. 5
the components of a! a" at each iteration until they are all in 1-aB, a.}. We now prove this intuition below using Lemma C.12, which proves the second part of the theorem.
We have g(a (0) = 2 also have , = (a and θ = {al X(a a²) E}. Applying Lemma C.12, we get
Xa (c) y"112, h(a(0) = ||a(0) 11, f(a()) = g(al)) + 3h(a(()) and s(a") = 3||aa ||1 + 2 We
|lalt)) - a= -a-
+ (tz - t1 aa max,,ssse [((((
((a) - f(a')) v 5
min
Vt2 2 t
(87)
2a(t2 t1)
Using [F(a())) O (3) Vt 2 t (Equation (84)) we get from Theorem C.13 that there exists C" > 0,
(f(al1)- - f(a")) S (1) + C'a B
|a a 11
min
to 2
(88)
fs St
at8
And it n * 0, i.e. E = 0,
(Ia - 1aai-llaall) S
7
||a!t)) - aa a²1
min Srrtt
+ 12 2 f1 +
(89)
2
a337
We now prove this result in a general case & + O. Let R(a): = (a aa) T G(a). We have from Lemma C.12,
[(aa))) + 3 (Ia( || ||^)|1)|
(o(al s(aa)) 3 (Ia -Ila-"a"" ||a*||] 11
min 4 S1St
<
min
4 SiSh
(90)
(f(al()) ([(a(i)-s(a7) f(a")) 442 24
min
Stst
So
B min
- -Ila*(1)) (Ia*")) s
min Gstts
((al))-((aa))
min g(a(0) - -
vta 2
(91)
StSt
Since
((a 1 .(a) x (0) a)1 - - 1 X((a a))l - U
1 - 1 E1
E >
(92)
We obtain
B min
(la*))) - -1la*111)
5 min
((aa)))((a*))
+ + 11 E1 v12 24
S15t
sSS5
(1 C'a3) B
CafB²
1/2 2 + *t(nn, since S /Cag
(93)
2
2
(1 (C C')a3)
12 2 2 2()
Lemma C.30. Given a > 0 and (1) R, let (++1) = (t)) - all (a for all t 2 1, where (a)) E Ola..
1. point a stationary for this dynamical system if and only if [a] S a.
37

---

Grokking Beyond the Euclidean Norm of Model Parameters
2. We have |a!0|| S a if and only ift >
aa
1
3. In particular, for h(a)= sign(a) Ya R, ifa" E Z, then (1 for all |
Proof. Let first consider the simple case (a) = =igg((a), so that ((++)) = alt) sign(a()))
If (+) E [0, a, -a, then ((+A)) 0 for all A > 0.
+ If a E (0, o.., then (++1) (t)) a E (-a, 0). and (t+2)) = =a (t+1)) + a= ====" =a E (0, a), and so on.
If (1) E (-a, 0), then (1+1) = (1) + a E (0, a), and (ll+2) = a1++1) - a = a E (-a; 0). and so on.
a > a (resp. ((4) < aa) it will be decreased (resp. increase) by a until ( E (0, a| (resp. a0)) E -la,,0), and we get back to the previous cases. In that case, |a(++1)| = |a!0| - a = |a - to S a t ++1 1
Now consider the general dynamic (+++1) = a0)) - a (a))) If (1) + 0 (the case a = 0 trivial), then the dynamic (t++1) (t) a (agg(a()) as long as |a 2 a, after which it will just oscillate in the ball {a, |a| S a} indefinitely. In fact, fixed point a must satisfy
(a): i.e. H(a) = 0. The only case where 0 a|a| is a = 0 or when it lies in the interval where the subgradient can be 0. However, for any such that |a| a, it is possible to choose (a) (for instance, H (a) a/a) /) such that = - (a), making a fixed point. Conversely, if la > a, then [ (a)) = 1 and la- - (a) = [|a] - al > 0, so a is not a fixed point.
Lemma C.31. Given a > 0 and (1) E R", let a (++1) = a
(al for all t 2 1, where (a) E 0|al|i-
1. A point a is stationary for this dynamical system if and only if ||aloo Saa
2. We have |lo0 S o if and only ift > [ l.
3. In particular, for h(a)- sign(a) Va E R", we have |a llo = K Kilaa aa E 2 for all
|
Proof. The proof is immediate by applying the Lemma C.30 coordinate-wise.
Lemma C.32. Let a,a" a" E R with a # O. If |a]
> |a*). then sign(a -a") aa) = sign(a))
Proof: We have |a*| < |a| + lal < a² < |a| = laa a" < 0 < |a| - aa. This implies a a" a² < for a < 0 and 0 0< a a" for 0 < a
Lemma C.33. Let f(0) g(0) + Bh(O) be a function from R" to R, where is a differentiable function and h is a sub-differentiable function. Consider the subgradient descent update f(++1) = B(t) - (Vg(0()) (3 (B(0)) with a fixed small step size >> 0, where H(0(0) E Oh(0). If ()()))))))l2 > [Vg((())) then the regularization term dominates the updates, causing the sequence (0(t) Jess to exhibit oscillatory behavior without convergence to a minimizer of f. The condition 3) ((())))) > [Vg(0(1))22 writes Byn > [[Vg((()) ll2 for &1 regularization, h((0) = =110|] V0 E ER"; and / /min((1, n22)= >> [V for for l, regularization, h(()- = |10| vO E Ra pa.x222
Proof Sketch. Given that ( 112 ss (Vg(0()) and H(0()) ~ H(0l1)) at the beginning of training, if (1)H(0)))]12 > Vg(( then (1(H()()))22 > [V)(0()) This inequality implies that the regularization term dominates the update, ((++1) xs 010) (O(E), with the influence of (g((()) becoming negligible. Consequently, the iterates do not converge to a stable minimizer of f. and the training and test error metrics oscillate, remaining above some suboptimal value.
For E, regularization, h(0)- = |0 v0 R", given that [IA(0())]2 & vn at the beginning of training, if B/nn >> [|V)(0()))l2 then that the update becomes dominated by the f Because H (((()) reflects the sign of ((t), the update effectively pushes the iterates in a direction that primarily depends on sign changes rather than the curvature of g (Lemma C.31). This leads to overshooting and sign flipping in each coordinate, resulting in oscillations.
For l, regularization, h(0) = ||@||= 4o E RB1 R=*x*",, the subgradient H(((t)) of ||0(")] satisfy (H() x
/min(ny. n2) at the beginning of training (full rank matrix), so ((())) 2 [| rannk (aan//ankk(H)))) (O(e)) R / /min(nn n22/ min(n1, n2) = / min(n1,n2)) If ///iinnn... n2) > (V)()))) then the update is dominated by the l. -term, making the iterates swing sharply depending on the current singular- vector configuration (Lemma C.42). LI
38

---

Grokking Beyond the Euclidean Norm of Model Parameters
C.5. Proof of Theorem 3.3
For vector u &R" and set S C [n], we let us = [u.]ies (RIS); and S [m]]]]]]
[n]],S.
Definition C.34 (Null Space Property). A matrix A E R=*xn is said to satisfy the null space property relative to a set S C [n] if ||us|] < ||uu||a for all non zero vector u E R" in ker A. It is said to satisfy the null space property of order 8 E N if it satisfies the null space property relative to any set S C [n] with |S| S 8.
If we add ||u||] on both side of ||u|||| < ||uو|| we obtain the equivalent formulation 2|[us||| ||| < ||u]|| i. On the other hand, by choosing S as an index set of s largest (in absolute value) entries of u and this time by adding | to both sides of the inequality, the null space property of order s reads ||u|| < 20 (u), where (u)) = inf Ivllo I|u - vllp. In fact, op (p(u) is achieve (but not only) V (u), with H₂ the hard thresholding operator (it keeps the s largest entries of u in absolute value, and set the remaining 0). We have ||u (u)ll ||ug lp with S the support of H, (u).
In theory, sparse recovery methods are designed to recover exactly sparse vectors. However, in more realistic scenarios, the vectors we aim to recover are not exactly sparse but can be well-approximated by sparse vectors. In other words, while the true signal may not have strictly zero entries outside a small support, most of its energy or information content is concentrated on a small number of components. This motivates the need for recovery guarantees that extend beyond exact sparsity. In such settings, we no longer expect perfect recovery of the original vector. Instead, we aim to reconstruct a vector such that the error between a and the true signal a" is controlled by how well the true signal can be approximated by an s-sparse vector. In other words, the reconstruction error should scale with the sparsity defect o, (a"), which is typically measured by the distance from the signal to the set of exactly s-sparse vectors. A reconstruction scheme that provides such guarantees is said to be stable with respect to the sparsity defect (Foucart & Rauhut, 2013).
Definition C.35 (Stable Null Space Property). A matrix A E Rm is said to satisfy the stable null space property with constant p (0, 1) relative to set S C [n] if [[ug]|a [1 < e|[اg| for all E ker A. It is said to satisfy the robust null space property of order 8 E N with constant p E (0, 1) if it satisfies the robust null space property with constant p relative to any set S C [n]] with |S| S 8.
Theorem C.36 (Theorem 4.14 (Foucart & Rauhut, 2013)). The matrix A R satisfies the stable null space property with constant p E (0, 1) relative to set S C [m] if and only if [|v ull1 s ++ ( ||u||| + 2|| for all vector u, V.V. E R" with Au Av.
i-p
If we apply this to the noiseless version of the problem of minimizing ||a]]] subject to | - y"ll2 S c, with a the solution return by, say basis pursuit, and 5 the support of (a), the we get ||a - a²I|, I 1(++) (a) since ||a*|| || S ||||]] and 1-p
||as||| = (a)
In practical applications, it is fundamentally unrealistic to assume that we can observe or measure a signal a' with perfect, infinite precision. All real-world measurements are subject to various sources of error, such as sensor inaccuracies, quantization effects, environmental noise, or other imperfections in the acquisition process. As a result, the measurement vector y that we actually obtain is not exactly equal to the ideal linear measurement Xa", but only an approximation of it, This deviation is typically modeled as an additive noise term so that we only know that the measurement error is bounded, lly" - Xa" ||2 = Ella scheme should not aim to recover a* exactly since doing so is impossible, but instead produce an estimate that is close to the true signal a", with an error that is controlled by the magnitude of the measurement error e. That is, small changes or perturbations in the observed measurements should only lead to small changes in the reconstructed signal. This desirable behavior is known as the robustness of the reconstruction scheme with respect to measurement error.
for some known estimated noise level c 0. In this noisy setting, an effective reconstruction
Definition C.37 (Robust Null Space Property). A matrix A E Rmxn is said to satisfy the robust null space property (with respect to || ||) with constant E (0, and - > 0 relative to set S C [n]] if ||us 1l < ell jluull 11 + A|l for all u E R". It is said to satisfy the robust null space property of order E N with constant p E (0, 1) and T > 0 if it satisfies the robust null space property with constant p and T relative to any set S C [m] with |S| S s.
Theorem C.38 (Theorem 4.20 (Foucart & Rauhut, 2013)). The matrix A E Rmxn satisfies the robust null space property
with constant p E (0, 1) and T > 0 relative to set S C [n] if and only l[ - ull1 S 1t (IvVll - ||ul]1 + 21[us|la) + 1-p Al for all vector u, v E R.
Theorem C.39. Assume the matrix X E RRn' satisfies the robust null space property with constant p E (0. 1) and r > 0 relative to the support a² Then, under the same condition as in Theorem C.29 on o, (3 and E; i.e. 0 < ax (X T X) < 2,
39

---

Grokking Beyond the Euclidean Norm of Model Parameters
0 < ByT < maxx (X (XX) X) and
vCaB with C O; there exists C' 0 Ouch that for all 0,
||a(0] - aa111 S Cin+ C20B + CallEl|2 = 12 2 + Al(n, 4),, A1( = ||a((.) a² 11
min 1
(94)
a.9m
with
P2 = max (I -
C 2(1 p))
C + C3 1+
2 max (X)
27
- p2) + p) + 8 Synommas (X)r
(95)
C2 =
2 1-p
1 1-p2 P2
1-p
2(1 - p2) p)
4r
C3 1-p p
Proof. Using Theorem C.38, we get |aa a <
(1 |a'll)) + T-p x - a²) We also have
ming (Ia(1)/2 (1 [a] I ++(C+C if and only if to + t2) (Theorem C.29). For 2 t1, have
[((a() (t) -a")l2 = ||aa() (1) - y + E
< |Xaa (t) -y"112 + ||E||2
2/ X 2-+2
0: |Xá 1--ly-l y21|2 ||E||2 (Equation (82))
(96)
1 P2
2/71
X 12-22
03 E|
since |Xá- y"1 = E (IN UU " < ||El12
(Equation (72))
1 P2
So,
min [|a!) 45t5
a²I|
1 + pn (C+ C')a3
2r
22 (2 (X)a3
= = 1224 22244 + 2/16112)
1 -
2
1- p
1- P2
(97)
=Cn + C20B + C3llEl12
- l 2 4 + (1(1114)
C.6. Proof of Theorem 3.4
Let (A) = X vec(A) for A E [Rm1Xng We have y X vec(A") + &, and want to minimize f(A) = ((() + BK(A)
using gradient descent with learning rate a > O, where 1(A): : ||. and ( = ( y"113 11 = { Xa
(X Xa* + xx a } Xa* + El. We write F(A) == G(A) (H{A) with vec G(A) = (a)(A) = X
[XT Xa" +X {) and H(A) E A|I = (UVT WW, | 12222 < 1, U. W = O, WV = 0} any subgradient of ||A ||x,
with A- = U under the compact SVD? 8, Suppose we start at some (1), Using F(O) = F(A()), ()), the subgradient
update rule is
(1+1) = (1) - -a((1) = (In ax T x) a ( +a +
vec T Xa" + x g) - a3 vec (H(A
14 >
(98)
As Section C.4, we let X U. VT under the compact SVD decomposition, with diag (g((TT)))et where 0. We assume the step size a satisfies We define Pp II x X p-p for all p > 0. We will show that for 3 small enough, the update first
r rank(X) and Fmax σi 2- ... 0x
0 2 0k+1 2 Jmin or > 0,+1
0
2
mmax
moves near the least square solution of the problem, á = vec Á (x x (V a" +E-AUTE) Later in training, (A) dominates the update, leading ||A (0)] A I|..
"The norm ||] not differentiable everywhere because the singular values of can non-differentiable at points where they have multiplicities (e.g., when the singular values are not distinct) 'For the experiments, used the polar factor (T(A) UVT. This is the gradient provided by automatic differentiation in many optimization libraries, like Pytorch.
40

---

Grokking Beyond the Euclidean Norm of Model Parameters
Theorem C.40. Assume the learning rate, the regularization coefficient and the noise satisfy 0 < < O' 0
Amar (X' X) and | S vCaB, C > O. Let P2 == mmax (In- - ax x...). There exist tr < 00 and constant C" (">0 0
(min(mi.ma)
such that:
20.5n!/2
and 9(A S gg(A)+ + 2na23oom max (X)
Vvec(A (1) -A]l2 - 1- P2
Vt 2 t1, with g(A) - 2 _lel22
(1 P1-)
(F(A(1)) - 5(A-)) - (1+ C'a:3) 3
+ 21 + At(n, 41), at(n.t): = ||A (1.) 1 (99)
Vn > 0,
min StSt
2
a:31
[|A5)], -I1A7))) 7 (C + C* aB
vn > 0. min StSt
- 12 22 (1)
2
Proof. First, we observe that if 33 is too high, the subgradient term (A) dominates early, and there is no convergence,
i.e. no memorization nor generalization. In fact, if 3 >
then the f,-term dominates the updates, causing the
m2))
sequence A (t) to exhibit oscillatory behavior without convergence to minimizer of f(A) = g())+ (Lemma C.33).
The memorization phase is similar to that of Theorem C.29, with n = 71 For t 2 4, g (A - < 2na 308 o (X)
(1-p2)²
O (2n max (X)) and || vec (A) S θ (28 (XT (x)) θ (3). So after time 41. the contribution of the gradient G to the update of A (0) is dominated by the fe-reullarizatoon term. Specifically for all t >> 41, the update rule
approximates A (++1)
(t)) (A.)( A(t) a.H (t) A)) up to error E E R21x==
x#2 of order ||E||
TayTank(A AUD)
o
(Lemma C.45). So, soon as the singular-value gap widens (i.e. min; (2(+00 > max; |
Omin (A(t) o/ (A+)-1
empirically typical after a warm-up phase), the approximation becomes tighter, even more so for small a:3. By Lemma C.42, this lead to (IA(() A1 Oaa3) for (and only for) t 2 t + ||A''' 12-2 As the error |E|| | of order O (a.)
o3
for iterations directly following (1, choosing time @ () also somehow counterbalances its cumulative effect. Equipped with this insight, we prove the exact delay below.
We have ((A (9)= X vec A()) y
12 ((A(A( = ||A f(A = ((A + (B(A (1) and f(A*) = 3/|A 11+
2 EI11 , = } and Gg = (A
X vec( - A' = E) Applying Lemma C.12, we get
S || vvec (A - -A*) 1 + (l2 t1)a² S1S12 |l vec (F(A(")) l
min 11 StSt2
- ([A(())-((A))
V12 2 4 (100)
2a(12 - (1)
Using
|| vec F(A (1) ]|| = (3) v4 t we get from Theorem C.13 that there exists C² > 0,
(S(A(0) ( (((7)) < (1+C'a 12 2 4
|| vec (A(t))
(1)) - A-) 11
min S1St
(101)
2
a:81
and
min (f(A())- - -((a)) - - min () 11E111
(IAA1)) - (1149).--1A711]
min f1 S1St
<
Vt2 21
4 45t522
+ C'a 8
1 1 E + to 2 tt(
(102)
2
(n+ + (C + (C++C")as) C')33)3
B
12 224 4 + (( since |X Ell2 S Vaa
since (1) - 111111 _ X ( A) U E N | E
Lemma C.41. Let A E Ra1 X22, We have || vec H) ||p < (π1n2)²/p
1/ for all H E A||||| and p > O.
Proof. Let H E A|.. Then |H|2222 S 1. So by the definition of the spectral (operator) norm, we have
|H] 12-22-22
|x]||
X max (H) < 1. Taking x e, m>)) the j-th standard basis vector in R"a we obtain H
supx/0
[He)") (12) ||2 5 1; which implied H, < [Hal 2| S 1. So ||
= D H,y)1) 1/p < (// 1/p.
41

---

Grokking Beyond the Euclidean Norm of Model Parameters
Lemma C.42. Given a > 0 and A E 182,, let A (L+1)- A - (A for all & 2 1, where H(A) E A||
1. A point A stationary for this dynamical system if and only if ||A 12-22 = maxx (A) < a.
2. [|A(()/22 122-2 < a if and only ift >
12
3. For all t
||A 1| x r= (a)
= {i
=
/ E 2) with of ......rr oi the singular values of A (
Proof. We start with the subgradient ((A)= UVT for A= UEVT, T, so that the update rule becomes
A ( A - -OU = (Σ() aln) v for all t 1
(103)
This equation also writes
Te41
A (4+1 -1-0))((
(++1)
U (++1)
v((+))
(++1)
(104)
Ft.
Σ (0
(t) aa)
10 (t) al sign(o; a)
1=1
This implies of (t+1) = 10,00 - al
Vi E [ri]. So starting at each のj decay at each step by a until of = oj E (0,
(1)
and start oscillating between of and og. It starts doing so when t > t ==
J. We take max, t,.
In general, we have ((A) UVT + W. W U,U W 0, WV 0. The dynamics becomes A (1+1)
U (E(0 al,,) V a ( with U W O, w() 0. So U ((+))())) g( ab, leading
again oz [++1] = 1o (0) al vi E [n]]
We show in Lemma C.32 that for a, a" E R. |a| > |a*| = sign(a a²) = sign(a). with sign(a) E alal. Now, considering the canonical nuclear-nore norm subgradient (AA) = UV E A||A||, for A = UTVT E R"11xng, we ask if
mmin (A) > muax (A") = H(A- - A* = H(A), and if not, how does H(A- - A*) deviate from (AA). Throughout, let A- = UEVT E R x of rank 2 1, and A - A UEVT R XM12 with U, Ú E RRIIx and V, V E R having
orthonormal columns. So (A) = UV and H(A A*) = T
A = 0 1
There exist matrices with minn (A) > o max (A*) for which H(A A²) + H(A). =
But
0
0
it easy to check that if minn (A) > maxx (A*) ) and A and A have the same singular directions, then H A') (A).
Lemma C.43. Let A UEVT E RRI xn2 If A UBV with mmax (B) < Tmin (A), then H(A- A*) = H(A).
Now we focus on bounding || [A A+) - [A)ll in a general setting. We introduce some notations, then Wedin's sin θ bound (Wedin, 1972), and our result. Let Sn = UU) C R" and 3 = sūa c R"1, The principal angles S 0, S ... S 0, S 7 2 between St. and 5r, are defined by cos 0; = a, U Ú ÚUÚ) Vi E [r]. where σs 2 ... 2 a, are the singular values. In fact, the classical (geometric) definition of the principal angles is given recursively by cos maxues2. ueš, u ū, cos 02 = max uex uiu, u " Each pair is taken orthogonal to the previous ones, so we end up with angles 0 .... < 0,, all in [0, 7/2). Write = Ux ū U with x, E IR, ||x| |*| 1. The first maximisation therefore reads 0 max =||=|||||= U
(U (U) The orthogonality constraints in the subsequent steps force (Xx: xx) = (U uj. U ūR)) to lie in the left-over singular subspaces of U Ü, with (u,, üj) the maximisers at step k. Inductively, one obtains cos 0; or (UTÚ) (i=1..... r) standard consequence of the SVD and the min-max (Courant-Fischer/Ky Fan) principle.
It is customary to collect the angles into a diagonal matrix θ, == diag(O)..... 0,). Applying sin(-) entry-wise gives sin θ₂ := diag (sin .,, sin 6,) One never needs to compute the angles explicitly because ||sin 12l1 = sin 0 = |IU U 122- where 0max = 0, is the largest principal angle, and Ú+ is an orthonormal basis of St. In fact, let B = U E Rexe and C = Ü- - E (r), We can write the SVD of B as B cos O,, ET, with
42

---

Grokking Beyond the Euclidean Norm of Model Parameters
E Rx orthogonal. This gives BB D O DD. Because of the orthogonality relation BBT + CCT (TÚT + Ú-Ū-TJU = U = l,. we have CCT I, - BB T = D(I, - (cos DT = ((sin DD) T.
So. under SVD, C D sin F with F ((==--)x((1-- orthogonal. This is related to the CS (cosine-sine) matrix decomposition of U [ U... Therefore. ||C||2-2 ||sin 12 = sin mmax.
Thus || sin 12 measures the worst-case misalignment between the two subspaces; it equals the operator norm of the difference of their orthogonal projections. In the following, we use sin ,, in place of || sin 0r.1|2-22 by abuse of notation. The following result can be derived from the Davis- Kahan sin θ theorem.
Lemma C.44 (Wedin's 's sin θ bound), Let E E RRIXEZ and set := := E. Let U,, U, E R span the leading r left singular subspaces of A and A, respectively, and let O₂ be the matrix of canonical angles between these sub- spaces. If |E 2- < or (A) (A) (with 0y (A) = O because rank (A) = =r), then sin Or == | 12 S
Omazl (E)
Plaxx (E) An analogous bound on sin On == V V 1l2-2 holds on the right singular side.
(A)
2 2/2rank(A)
Lemma C.45. m (A) > maxe (A*), then [A A A - (A)))le S
(A) T (A')-I'
Proof. Let (A). Write ((A) UVT and H(A A²) ÜVT, with U, V, U, V having orthonormal columns.
||U - - = U i | - 2r (ÚTU) = 2r - 2 Σ
cos 0g since | = ر = and o; (ūTV) - cos 8;
1=1
sin 0g since 1 - cos 02 = sin² 0:
= 2 2 (1 cos 0,) 2
S sin" 0, for 02 E [0, x/2
1 cos 0;
i=1
S 2sin m = 2sin 22)²
mmax (A*)
2
S 2r
(Lemma C.44 with E = AA")
Uiin (A) - O maxx (A*)
(105)
Hence
UUVT UV I I(Ū U ||U( V)T(F |F
2v 2ro max ( (A*) minn (A) mmax (A+)
(106)
< | U + / -VVv -
C.7. Proof of Theorem 3.5
Definition C.46 (Stable Rank Null Space Property). A linear measurement map F RR1XX1 + RT is said to satisfy the stable rank null space property of order r with constant E (0,1) if for all A E ker F \ (0), the singular values of U satisfy Ei=r (A) S Eistt min(()).e)) (A).
Theorem C.47 (Exercises 4.19 in Foucart & Rauhut (2013). The linear measurement map F [Rn:Xn2 - RT satisfies
the stable rank null space property of order with constant & E (0, if and only if |B All+ S 1- B | |||| +
2 Exert m22) oz (A)) for all vector A, E R Xm2 with F(A) = F(B).
Definition C.48 (Robust Rank Null Space Property). A linear measurement map F R
- Rm is said to satisfy the robust rank null space property of order (with respect || II) with constants & E (0, 1) and 7 >> 0 if, for all A E Ruix. the singular values of A satisfy Et= (A) S Eminkn ae(A) + +T]]F
:R X
Theorem C.49 (Fxercises 4.19 in Foucart & Rauhut (2013)). The linear measurement map F R X - Rm satisfies the robust rank null space property of order (with respect to || I|) with constants p E (0,1) and T > 0 if and only if
BB- All+ < 1+e 1-p (I ||A||| + Ex
min((n..n)) ,22) (A))A) + 2" [I A)) for all vector A, E R
Theorem C.50. Assume the linear measurement map F. satisfies the robust rank null space property of order with
constants p E (0, and r > 0, i.e for all A E Ra X01 Σi (e) S <p Emira (2)))
ay(A) + [[ee(())(())l Then,
43

---

Grokking Beyond the Euclidean Norm of Model Parameters
under the same condition as in Theorem C.40 on a, and &; i.e. < 00 X) <, < ( m2)
mmax (X (X) and || Ell2 - vCa3, C' > O; there exists C² > such that for all 1] > 0,
41(1, = ||A A 1 (107)
min S
(A A + 1- S Cin C20B + C3l1E112 = 12 2 4 + 4),
o.3n
with P2 = 0 max (In - -ax axx x).
1 1+1 p
C C+ C'
2 no max (X) 2r
4r
== = 2(1 pp)
C2= = 2 1 1-p
+
and C3 =
(108)
1 P2
1-p
1 1-6
Proof. We have | A"I.. S
1+e - |1)))
T-p X vec (A A (Theorem C.49). We also have (C+C)aa if and only if l2 2 h + A1(1,1) (Theorem C.40). For t 2 f1, we
ming Sesta (II (0) ||- - ||7lla) S
2
have
X vec(A - - 1 X vCc (8)
y" 11a + |E112
T/
X
2-+2
a.3 + +| y'll2 + ||El|2 (Equation (82))
1 P2
(109)
2
X
22
0.3 + +|11E||2 (Fquation (72))
P2
Combining these gives the desired result.
C.8. Proof of Theorem 3.6
We want to minimize f(a) g(a) + Bh(a) using gradient descent with learning rate a, where g(a) | y'113
and h(a) ||all/ Let Q := X X -la. We have
f(a) = 2 Iy(aa - -y"11 2 1lall
= a
Qa TaQ
Xa" +x xx
a Xa" + +112
(110)
2
and
F(a) = vaf(a) = x XT (y- 3a Qa Qa-ax Xa + x)
Xa*
(11)
The subgradient update rule is
a (1-1) aa) (F(a = (In - aQ) + (X (X Xa" + +X
(112)
Let X = V under the SVD decomposition, with Σ diag(ox)<Ein)> where r = rank (X) and mmax 0. If the step size satisfies 0 < a F The update a!)) converge to the least square solution å = Q+X y", but this solution can not give rise to generalization when N < n.
2
Tk Tk+1 min o, 0r+1
Theorem C.51. For all p > 0, let define Pp == I (x x + 31m) 1I Assume the learning rate satisfies 0 <
< mman (X XyF3 Then ||a!() - allp S oz (a(()) - allp for all 2 1. On the other hand, for N < n, || a 13 >
(10, VVT T )a* 112. In particular, af has a nonzero component orthogonal to the column space V, then à cannot
perfectly generalize to a*.
Proof. In Lemma C.52, we show that [|a(+1) - al1, S <filla!") - al1p AI 2 0. So for L 2 In (1/ - -(ll)) In Pp. we have (++1) - allp < Consider the regularized least-squares estimator á = -xx + 31,) x y* V - 31) E²U y*. We have VVT à. i.e. â E Col(V). Let decompose a" into two orthogonal com- ponents; a' = a + a1, where aj == VVTa" E Col(V) and a = (0, VV )a Col(V)). Since Col(V), VV (a a = â ---a al and VV a = 0 by orthogonality. Thus, we can express the error as à a' (al + a (a - a a Because â all Col(V) and a lies in the orthogonal complement of Col(V), these two vectors are orthogonal. Hence, ||á a" a'I = la 11l1 + ||a 113 2 |lli 1 = |l(1, vVT)a*|13 LI
44

---

Grokking Beyond the Euclidean Norm of Model Parameters
Lemma C.52. If (0.
Tmax sm then F(a!0) + as + oc; where F(a) 7 à +
x (h--(x + ) x+ Bln)) c= à + (0n, - VV c v R". Also, [|a - ållp S P/lla -ålpz vt.
Proof. The solutions of F(a) = 0 are
X X B1n) a= y
X X + )lm) X X x'xa" & VEV a + UV
(113)
- (X x + (1m)
x y In (x x 31m) (x + (l1,)) c a (In VVT)c Ve E
Let E = - a (Σ + +11;; A = In a (x- + 31m)) = V.V and = ax a Xx Xa" + x
VEV a' V E) that
f-1
a ((+)) = Aa (1) + A +
Σ 1=0
A'
A'a +00 AO- (1- A') (I-A")
w
(114)
Ast-00, E A' _ 0. We have
t-1 Σ 0
1-1 A'= -n + Σ
VL'VT = 1n VV
Σ
V
1m1
1 (1-0% V
0/ 1 )
ln VV T + V diag
Σ o
=I = - VV + V diig
(115)
1=0
= In - VV + V (1-
[1 (1-50) VT
-4 12 VV + v (1 (1-E)
-1 V = ln -VV + (E V as t +
So, as + 00,
a(++1)=
Σ A² )
8
w =a 1.. VV + 1 (E+ - v+) (VEVTa a VEAUTT)
a
1=0
= VE + 30) -1 'v_ (VEVTa a' + VE UT 5) + (I, VV c with c= w
(116)
= VV 30) 1 vv
(VEVTa a" -V =
We have Aâtc- å, a (t (1-1) A(a â) A'(a(() -)).
(a)))) which implies ||a [t(++))=aa) S ||A'11p (p--)paa((()))
D. Other Iterative Method for l, and l, Minimization
D.1. Sparse Recovery
For the experiments of this section, we use (n, a. B) = (10 11-0.00--- 1000,.001' 10 1. We solve the sparse recovery problem using the projected subgradient (Figure 12) and the soft-thresholding algorithm (Figure 13). We observe a grokking like pattern similar to the subgradient case. One training step is enough for the projected subgradient to get zero training error. This further shows that generalization is driven by a and, more importantly, 3.
Projected subgradient To ensure memorization, we can use the projected subgradient for problem of minimizing ||l]l subject to the constraint Xa = y", where at each step the update (using now just (a) as gradient, not the whole F(a)) projected onto the constraint set. In our case, the update write (++1) IT (a (a()) with
(Ta)- = a T (XXT) (Xa- y*) the projection of a on the set (a. Xa = y*). We can show that the e optimal
45

---

Grokking Beyond the Euclidean Norm of Model Parameters
gap of this method enjoys the same bound ( given above (Theorem C.29) for the non-projected case. We have milis (|la(1//1 |aalll)) II]) (a()) (12))(tt)t))))))))))))))
ad
The proof the following:
25a (1)
22 2
0 < |a (t+1)
a²112 ||| a (tz) aB. (1 (2) a
|a (ts) a" - -33 ((a (1a()))
(aa a" 1 20 (12)
aa") Haa"H((( + 8 (
(117)
(a())) a 1 23a
||a ||a(ta)|| ||a'|11 fa H ا1
(H (by the definition of H)
(Ia ||a"111) + 0 42
(a(tt) 2²1 23a
H(a
This implies
1/
(Ia(1) 111 - ||a"|1 111 I ||a(t)) -a 11 32a 2 Σ
23
a
min SSS5/2
1=t1
t=4
(118)
(|a|0|/| - ||a*)] (|a*11) ||a(t)) a²1 32 (1) maxe, S15ta (1
min siSt
2,30x11 41)
s=1 1
s=5
100
100
0.8
80
0.8
80
70
70
Error
0.6
(1
1l2 / 11y ll2
Error
0.6
1|a
-y
112
1y
112
0.4
||att
12
a|a ll2
50
0.4
112
||a
40
30
0.2
20
20
0.0
0.0
10"
10a
102
10a
10
10a
10
109
10a
102
10a
10a
10a
10
Steps (t)
Steps (t)
5=10
s=15
100
1.0
100
80
0.8
80
0.8
70
Error
0.6
||a
1l2
/
ll2
Error
0.6
-y 112
11y llz
0.4
(|
112
|aa
50
0.4
||a
$12
||a l|2
40
0.2
0.2
20
20
0.0
0.0
100
10a
10c
10a
10s
10a
10a
102
10a
10a
10
106
Steps (t)
Steps (t)
Figure 12: Sparse recovery with the projected subgradient descent method.
Proximal Gradient Descent We have a - aG((a) = arg min. g(a) + (c ((c-a) GGa) + L ||c - alž. So
2n
1
a - aF(aa ss arg argming(a)) + (c - a) G(a) + - |c -c-al13 + B||
2a
arg min 2a
||c (a aG(a)) 1 + 3|||||1
= II, (a aG(a))
46

---

Grokking Beyond the Euclidean Norm of Model Parameters
with aa the proximal mapping for + |||c|| (a) arg mine 1 -c || arg minl 1 - Sag(aa (a), where (y)a) (a) max([al 7, 0) the soft-thresholding operator".
a; - if 7
S (a),
0 if 7 a 5
a;
+ if < ~-
The final form of the update, known the Iterative soft-thresholding algorithm (ISTA) (Daubechies et al., 2003), is then X 2 max (X X) be the Lipschitz constant for G. If S 1/L, then minegr(((a1)) ((a)) ||"">="'' llz
a -Sa3 (a(*) (G() At > 1. Let L
0 for the ISTA. We applied standard bound on proximal 20l T
gradient descent (Tibshirani, 2015) for a function of the form f g h: R R. Such result state that the proximal gradient descent with fixed step size 1/L satisfies (a( f*) (|ا -ل"ام when is convex, differentiable, do R", V is Lipschitz continuous with constant L > 0; and h is convex and its proximal map II can be evaluated.
200
s=1
s=5
10a
10a
100
80
10-1
80
10-1
Error
| m
1l2
Byy ll2
Error
||a
1l / 1ly 11
10-2
||att
1l
||a
ll2
50
||
112
la
llz
50
10-2
10-3
30
30
20
20
10-3
10
105
10
10
102
10
10
Steps (t)
Steps (t)
s=10
s= 15
10a
100
10a
100
10-1
60
10
60
Error
(aat
1l / Ily 1l2
Error
IIXa
ll2 / Ily ll2
10-2
10
112
|aa
llz
50
||
lll
||a
llz
50
10
10-3
30
30
10-4
20
10-4
20
10"
102
10
105
106
10
101
102
10
10
10a
10
Steps (t)
Steps (t)
Figure 13: Sparse recovery with the soft-thresholding algorithm (ISTA)
D.2. Low Rank Matrix Factorization
For the experiments of this section we use (n1, n2, r, N,, S, a, 3) (10, 20, 2, 70, 10 10,10-110 1,10--4, for matrix sensing, Figure 14 shows the results for the projected subgradient and Figure 15 shows the results for the proximal gradient method.
Projected subgradient At each step, the update (using now just (H gradient) is projected onto the constraint set. In our case, the update write A (t+1) (A (A a:H((A (A))) (9)) with IT the projection the set {A, vec A = =y").
Proximal Gradient Descent We have A aF(A) 11, (A aG(A)) where 1I, is the proximal mapping for B + B||,, (A) arg ming 2n | A + B|| S.g(A) (A) with (A) U Umxx{E 7, 0)VT the soft- thresholding operator for A UEV
under SVD, where max(E 7, 0).z d.j max(E;; 7, 7). The final form of the update is then A (t+1) Sae (A (4) aG(A(t)) Vt > 1.
"On C. the soft-thresholding operator (a) = sign(a) @ max la| - 7, 0) only shrinks magnitude and keeps the phase fixed.
47

---

Grokking Beyond the Euclidean Norm of Model Parameters
10
Singular values
Dimensions (n)
15
10
0
( - " l1x 1l2
||||
10
1|A10 -A" 11= / IIA l|F
100
10-11
||A|
10
||A I1-
AF
10
10
10
10
10a
10
10
10a
10
10
10
10a
10
10
Steps (t)
Steps
(t)
Figure 14: Matrix sensing with the projected subgradient method.
109
Singular values
Dimensionn (m)
") >/1bb
1b
Error
10
1 All ||
10
II -y -y"llr
11''11z
|A||
10-9
10-3
IA - 11- 117
100
I|A")]-
10
10-1
IIA I|--
10
10
10
100
10
10
10- 10
10*
Steps (t)
Figure 15: Matrix sensing with the Proximal Gradient Descent method.
E. Grokking Without Understanding
Consider the sparse recovery problem (the explanation below also holds for matrix factorization). We start the optimization at (1) iid (0 1/n) with (2 O the initialization scale. With small initialization, & regularization is sufficient for generalization to happen, provided N is large enough and & regularization is not very large. If the scale at initialization is large, then adding & regularization is necessary to generalize, but is it sufficient? That is, can we generalize to the problem studied here with only le regularization?
shown in Section C.8, the answer to this question is no. But what we want to illustrate here is a phenomenon that contradicts previous art (Liu al., 2023a; Lyu et al., 2023), namely that in the over-parametrized regime (N < in our case), large initialization and non-zero weight decay do not always lead to grokking. What happens is that, because of the large initialization, a more or less abrupt transition is observed in the generalization error during training, corresponding
Ener
400
1
fu'll
200
Stops (2)
Stops
Figure 16: Relative errors and (a with large initialization scale Ç 10 and small weights decay = 10-5, Without visualization of the error on logarithmic scale (left), it looks like grokking has occurred, whereas this is not the case (right).
transition in the & norm of the model parameters. But this not be called grokking because the model only converges to a sub- optimal solution. What is more, this transition appears even if the problem posed admits no solution, e.g., sparse recovery or matrix completion with a number N examples far below the theoretical limit required for the solution to the problem posed to be the optimal solution (by any method whatsoever). This transition appears abrupt just because the training error large at the beginning training, since the model's 's outputs are large. When its l2 norm becomes small, its outputs also become small, leading to transition in error. In Figure 16, without visualization of the error on logarithmic scale, it looks like grokking has occurred, whereas this is not the case. For this figure use (n, 8, N, a:) (100, 5, 30, 10-1).
We call this phenomenon "grokking without understanding" like Levi et al. (2024) who illustrated it in the case of linear classification. They show that the sharp increase in generalization accuracy may often not imply transition from "memorization" to "undeestanding" but can an artifact of the accuracy measure. But in our case, we are not using any significant scale at initialization (we focus on 0 S Ç & 10 10- and are not dealing with the generalization measure problem since our test error directly the recovery error in the function space, not the accuracy.
48

---

Grokking Beyond the Euclidean Norm of Model Parameters
We hypothesize that the interplay between large initialization and small non-zero weight decay that leads to grokking as predicted (provably) by Lyu al. (2023) does not hold in our setting because our model violates they Assumption 3.2. Let Va (x) = aa x denote our model.
Assumption 3./ (Lyu et al., 2023): For all x e R", the function a + (x() is 1-homogeneous since yea(x) = ya for all c > 0.
Assumption 3.2 (L et al., 2023): for Ç = 0, yacn (x) = 0 for all x (there is generalization in this case with {1), but if > 0 (for instance Ç large), this is (almost surely) no longer true. So, this assumption is violated (with high probability).
Assumption 3.8 (Lyu et al., 2023): The NTK (Neural Tangent Kernel) features of training samples (V (VaHa(Xi)) HEIN are linearly independent (almost surely) at initialisation = a In fact, VaHa(x) = x vx. In the over wver-parametrieed regime N < n, If M E RN** has entries iid from a normal distribution, then the NTK features (Xr) are linearly independent with high probability (because the rank of X is N with high probability), so this assumption is verified.
F. Implicit Bias of the Depth
F.1 Deep Sparse Recovery
For sparse recovery, let now use the parameterization a = _f E R", with A E R4xm, This corresponds to linear network with L layers, where each hidden layer has the parameter diag(A, E R=x=_ with this, increasing L leads to overparamettrizztioo without altering the expressiveness of the function class a - (x) = x a, since the model remains linear with respect to the input x. Unlike the shallow case (L- = 1), there is no need for l regularization to generalize when 2 2 (and the initialization scale is small), as the experiments of this section suggest. With depth, the update for the whole iteration (which is now replaced by a product of matrices and a vector) similar to the shallow case, but with preconditioner in front of the gradient. This preconditioner makes it possible recover the sparse signal without any regularization.
We let h(A)) = 1 All2l so that H = A,, Also g(a) = 1/2 )lly(a) - y f Xa (x T Xa' + + <) a
| + E E12 Let G(a)= = @g(a)
X x(a a²) X E. The gradient for each A, is G(A.) := 0g(a) DA,
8a 0g(a) DA, Ba
da
diag (Oryt (a). We start the optimization at A (1) add for each A; is
~
(N(00, 1/n) with Ç 2 0 the initialization scale. The update rule
A (++)) = A -
- a diag (Oz ( G(a
(119)
Without ovaparametrizationn (L 1), the update is unconditioned and progresses uniformly in all directions. So without f1-regularization, there is no mechanism to enforce sparsity, and perfect recovery of a" is impossible. For I, = 2, let C: == A C A + A2 @ Ag.
a (++1) A ((+1)
② A (++1)
(1 - A (0) - a diag(A,)) (1) G(a @ (11 - a3)) a diaig G(a())) ((a)))
(1 a0:8)2200 a(1 a:3) diag A (4) の 1 + A の Al")G(a")) + a² diag(a ((a (120)
= (1 - (l-a3)2a (0 a(1 - o G(a ()) + a²a (0 o ())
(1- 2a (1) ( o G(a (1) for 0 0
and
c(++1) A A C A (1-1)
(1+1)
+ A2 3 A22 (++1)
(1+1)
(1 aB)2 (1)
c A 20 aB) diag(A A (1) o A (t) G(a a² diag(A
(9) (1) ((a0022
A
+ (1 - -1-a,))a a @ A
(1) - -aa a:3) diag A @ ()a + a² diag A @ A
(t) (1) ((a (121)
= (1 -03)) 4a(1 ") (1) 3 (( + (2 @ ())2
= (1 22 (1) - 4oa (1) o G(a for +
49

---

Grokking Beyond the Euclidean Norm of Model Parameters
The depth adds the preconditioning p(() (1 - a.3) ( in front of the update for a. This preconditioning mechanism seems to implicitly favor sparsity and, thus, a perfect recovery after memorization. In fact, when c, goes to zero (which is the case when a: is also small), the update becomes a (*+1) ~ (1 2a (*), and thus push a (+++) to at a geometric rate (t)
of O(1 2o.3). Otherwise, cf (large) will amplify the gradient so that of G(a(()); ); dominates the update, which pushes (* towards a" (as the gradient G(a(!) points towards a small error a (f) a" direction, particularly for full rank X and high signal to ratio regime).
With the deep parametrization above, we solve the sparse recovery problem using the subgradient descent method with (n,,,, s, a, 3)
(102, 5, 10 10-1.00 0), for different values of N and L E {1, 4] (Figure 17 and 18). We use small initialization scale Ç 10-6 for L = 1 and < 1 for L > 1. Here, initializing A too close to the origin (initialization scale Ç < 0) leads to not change during training. The model is able to recover the true signal a", and the generalization delay becomes extremely small (compared to the shallow case with non-zero ε coefficient) for I, = 2 and disappears (ungrokking) for L > 2. As L becomes larger, the phase transition to generalization becomes extremely abrupt. The loss decreases in a staircase fashion, with more or less long plateaus of suboptimal generalization error during training. This type of behavior is generally observed in the optimization of Soft Committee Machines (Bichl & Schwarze, 1995; Saad & Solla, 1995b;a; 1996; Engel & Broeck, 2001; Aubin et al., 2018; Goldt et al., 2020), which are two-layer linear or non- non-linear teacher-student systems, with the output layer of the student fixed to that of the teacher during training.
Additionally, for fixed number N of measures, the test error decreases with L, indicating that depth aids in finding the signal with smaller number of measures, albeit at the expense of longer training time. So, the depth seems to have the same effect on generalization as 4 This is in accord with the result of Arora et al. (2018) in the context of matrix factorization. They show that introducing depth effectively turns gradient descent into a shallow (single-layer) training process equipped with a built-in preconditioning mechanism. This mechanism biases updates toward directions already explored by the optimization, serving as an acceleration technique that fuses momentum with adaptive step sizes. Furthermore, they demonstrate that depth-based overparameterization can substantially speed up training, even in straightforward convex tasks like linear regression with L, loss, p > 2.
Note that for L 2 2, using large-scale initialization and a small but non-zere &g regularization 3 results in grokking, unlike the case of L = 1 that gives the "grokking without understanding" phenomenon. In this regime, when L increases, the number of steps required for the model to move from memorization to generalization is reduced (grokking acceleration), and the generalization error at the end of training is considerably lower. Lyu et al. (2023) used a similar setup to show that an interplay between large initialization and small nonzero weights decay gives rise to grokking with the diagonal linear network y(x)= (u - voL) x in the context of binary classification, but there did not study the impact of L on the generalization delay, and instead focus on characterizing how sharp is the transition from memorization to generalization as a function of the initialization scale and the weight decay coefficient, and how long it takes for this transition to occurs. This diagonal linear network is also often used for sparse recovery problems (Vavskevivcius et al., 2019), but the focus is generally on its ability to recover the optimal solution, not grokking.
F.2 Deep Matrix Factorization
For matrix factorization, let use the parameterization = II As, with A E RR1xd, Ar E RRdx,, and A, E daxd for all E (1, L). This corresponds to a linear network with T, layers, where each hidden layer has the parameter with this, increasing L leads to overparameterizztion without altering the expressiveness of the function class A- FAFA(x (x) = x vec A, since the model remains linear with respect to the input x. Like in compressed sensing, there is no need for E. (or any other form of regularization) to generalize when L 2 2 (and the initialization scale is small), unlike the shallow case (T- = 1). This is an observation already made and proven in previous art (Gunasekar et al., 2017; Arora et al., 2019; Gidel et al., 2019; Gissin et al., 2019; Razin & Cohen, 2020; Li et al., 2020). Gunasekar et al. (2017); Arora et al. (2019) show that increasing implicitly biases A toward a low-rank solution, which oftentimes leads to more accurate recovery for sufficiently large N. In fact, with depth, the update for the whole iterate is similar to the shallow case but with a preconditioner in front of the gradient (like in section F.1). This preconditioner makes it possible recover the low-rank matrix without any regularization and with fewer samples than in the shallow case (Arora et al., 2018: 80 2019). It is also shown specifically for this problem that initializing the model very far from the origin and using a small (but non-zero) weight decay leads to grokking (Lyu et al., 2023), i.e., the model first memorizes the observed entries, then after a long training period, converges to the sought matrices provided the number of such observe entries is large enough.
50

---

Grokking Beyond the Euclidean Norm of Model Parameters
N=10
N=20
1.0
1.0
0.8
0.8
Error
0.6
||aa
y 1l2
lly lla
Error
0.6
X|
1l
1ly l12
0.4
1|at0
1l2
a 1l2
0.4
||
112
||a
l1a
0.2
0.2
0.0
0.0
10c
10
10
10
10a
10
10
10
10a
10
10"
Steps (t)
Steps (t)
N=30
N=40
1.0
1.0
Error
||Xa
y 1|2
Ily ll2
Error
0.6
||Xa
y 1l2
1y 11a
||a
ll2 / a 1l2
0.4
||
112
||
112
0.2
0.2
0.0
0.0
10c
10
10a
10a
10
10a
10
10
10a
10a
10a
10a
10a
10
Steps
(t)
Steps
(t)
N=50
N=60
1.0
1.0
0.8
0.8
Error
0.6
||aa
y 112
Byy lk
0.6
||a 10
y 112 11 11z
0.4
||a!!
|l>
0.4
||a
11
/ I|a 11>
0.2
0.2
0.0
0.0
10
10
102
10
10c
103
10
10
10
10a
10a
10a
10a
10
Steps (t)
Steps
(t)
N=70
N=80
1.0
1.0
0.8
0.8
0.6
||
ll
lly lla
00
0.6
1 "
y 112
lla
0.4
||at0
11>
||a
11b
0.4
1aa"1
11
/ f|a 11
0.2
0.2
10c
10
10a
10
10
102
10a
10a
10a
Steps (t)
Steps
(t)
N0
N=100
1.0
1.0
0.8
0.8
Error
||Xa
y 1|2
Ily lla
x|
y
lly 112
11
ll2 /
a|a lla
||a
a lla / aa 11a
0.2
0.2
0.0
0.0
10c
10'
102
10a
10c
10a
10
10a
10a
102
10a
10a
10a
10
Steps (t)
Steps (t)
Figure 17: Training and recovery error as a function of the number of samples N and the depth L.
51

---

Grokking Beyond the Euclidean Norm of Model Parameters
10"
1b
11:
10
lla 112
10-
112 /
11
1004
10
10
165
10
1.1
27 101
153
10
Figure 18: Training and error [aa(t)) y 112/ [|la//lly"|l 2|2 and recovery error ||a(()) a //2/ ||*||2 (along with i and f2, the memorization and the generalization step) function of the number of sample N and the depth L. The growth (as function of N) in the test error for L is simply due the fact that did not optimize long enough for it to decrease.
G. Amplifying Grokking through Data Selection
G.1. Sparse Recovery
Definition G.1 (Restricted Isometry Property (RIP) and Restricted Isometric Constant (RIC)). Let A Rmxn and (8, (,) E [n] x (0, 1). The matrix A is said to satisfy the (s,6,)-RIP if 8,) < |Ax| < (1 + 1x)1 for all s-sparse vector x E R" (i.e. ||x||o < s). This is equivalent to saying that for every J C [n] with |J| =
(1 (1-6,) [x12 < | < (1 +1+ 5.) 5 Vx E RR where the submatrix A E Rmxs of A is build by selecting the columns index in J. This condition is also equivalent to the statement ||ATAA,) 0g 1212-22 8g, which is finally equivalent to AA., E 8,, 1 5,] for all eigenvalues eigenvalue À (AT,
of J.J.A,,J We say that A satisfies s s-RIP it satisfies (s, R.)-RIP with some 5, E (0, 1). The s-RIC of A is defined as the infimum (A(A) of all possible ô, such that A E Rmxx satisfy the (s, & RII. So, for all VJ C [n] with [J| = 8, the condition number of A.,. is bounds from (A)
above by 1+8,(A) I-6.1A) the one of A.,. by
I-)
We say that a matrix A satisfies the RIP if (,(A) is small for reasonably large All the above definitions extend to any linear map f IR" + R Note that (A) S d (A) for all A E Rmxx and E [n]].
Let Rmxx x + Ra be a linear map and (r, 8,) E [n] x (0, f is said to satisfy (r, 6 for all rank-r matrices
X Rmxn (1 8,) || < | S (1 8,) XII We say that satisfies r-RIP F satisfies (r, 8IR-RIP with
some 8, E (0, 1), and the r-RIC of F defined as the infimum 5, (F) of all possible 6, such that F satisfy the (r, RRIP
Definition G.2 (Coherence). The coherence between two matrices A E RTXm and B Rax= is p(A.B) matrices or vectors are. Specifically, measures how much overlap there is between the columns A and B. High coherence means they are similar or aligned, and low coherence (or incoherence) means they are very different. Incoherence essentially the opposite of coherence. It refers low overlap or low similarity between the columns of A and B.
[ Bl [a [[]. maxig[m] A.. B. Coherence measures how similar or aligned
[.A...B..)I
The mutual coherence of a matrix
E RMXR
(AA) = max
AA A |
[AT (AT [[
max((.))(m/ x(nlajj
If the coherence is small, then the columns of A are almost mutually orthogo-
52

---

Grokking Beyond the Euclidean Norm of Model Parameters
nal. A small coherence is desired in order to have good sparse recovery properties. We also have the 1-coherence sp(A).
I A..)) Jn s) = maxie[m] maxss/MAAes Eyel A.I.ll....
For a matrix A E [R"x" with unit norm columns, (u(A) 2 (Rauhut, 2010). In high dimensional space (large n),, ((A) 2 nn-I achieve when A is an equiangular tight frame.
a-m m m(n-1]
and pa (A,8) 2
m(n whenever & S Vn
n M
becomes (A
2 V
: this lower bound is
Proposition G.3 (Relation between f and 8). For matrix A E Rmxxn with unit norm columns, n( = (A).
(A s) = max |IAT, A. - 0111-1, and (A( S (1 1) S (s- ())(A) (Rauhut, 2010).
G.1.I. THE PROBLEM
Given the sparse basis (or dictionary) & E RR*x,, the mcasurement matrix M E RNX" and the measures y" Fb- (M) + E = Mb* + & E RN; we aim to solve (Po) Minimize ||a|lo s.t. [Ma y I c, and more pre- cisely, its convex relaxation (Pi) Minimize ||a|] s.t. (M++) - y" 1/ S e. This problem has been well studied in the signal processing literature under the name Basis Pursuit. It is well known that under certain conditions surement matrix M (e.g., coherence with respect to 2) and the sparsity of b' in D sufficiently sparse solutions of (Pi) are also solutions of (Po) (Donoho & Elad, 2003; Candes et al., 2006). Many lower bounds on the number of measures N guaranteeing ||a - a²I12 S with high probability have also been derived. Such lower bounds generally have the form N= 22 8 8 (s log C2 (n/s) + log 1/n) (Rauhut, 2010). where ó capture the Restricted Isometry Property (Def- inition G.1) of X M4 (X) S 5) and is also related to the coherence (Definition G.2) of M with respect to (Proposition G.3), 1) is the percentage of error (i.e. N guaranteed recovery with probability at least
the mea-
71, C₁ > 0 and C2 0 are constants. Observe that in the noiseless setting, we want such that Xa Xa*, that is E a" + (X). Donoho (2006b;a) show that the nullspace (a, Xa = 0} has very special structure for certain X (e.g. incoherent with any orthonormal basis): when a" is sparse, the only element in the affine subspace a² + Null(X) that can have a small & norm is a' itself
We will assume for simplicity that & is an orthonormal matrix, T = I,. It is common in sparse coding theory to consider & E Ruxm as a dictionary with m columns referred to as atoms (square, n m; undercomplete, n m: or overcomplete, n < m); and saying b" is sparse means it can be written as a linear combination of few of such atoms. But here, we assume for simplicity that we have b" taa with a" E Rm and 4 Rnxm a set of m I linearly independent vectors (its column). Let h- e (n-(m--) be the orthogonal complement of t in R", VT = i Φ-1 E RRax,,
:= V/ (v v) 1/2 E Rax= the orthonormal version of V, and ä" = (v T v) 1/2
Ta*)
We still have b* = .āa. with with
0
|ā" ||a" |lo since v T is diagonal. So, assuming + orthonormal is without loss of generality.
Example G.I. For the Fourier basis vnl ja e 2- we have in (0. s) = sμ(() s/ V (Rauhut, 2010). Each column
in this basis vector corresponds to a specific frequency. For a signal b*, if only a few frequency components contribute significantly to b", then a" = tb -b"bb, the Fourier transform of b", will be sparse.
G.1.2. THE CONTROLS PARAMETERS
The incoherence between the measurement vectors (line of M) and the sparse basis (column of ()) crucial for successfully recovering b" (or equivalently a", the sparse representation). If M is incoherent with t, each measurement captures distinet "view" of b", reducing redundancy. This diversity of information allows for the successful reconstruction of a* even with fewer measurements (e.g., below the Nyquist rate for signals). Achieving low coherence (high incoherence) can be done by designing M to be a random matrix (c.g., Sub-Gaussian like Gaussian or Bernoulli matrices). Such random matrices are, with high probability, incoherent with any fixed orthonormal basis, as Theorems G.4 and G.5 show.
Theorem G.4. Le m S and & E Ruxm with क क- = Im For any N 2 1, C1 > 0 and C2 2 1; the matrix M E R Vin(aN) with probability least 1/ 1/(nN)o03-
with Ci M itd N (0, 1) satisfies p(MT, 4) S 2C2
Proof. Let o- = n -i, C > O. We have M rid oO so [M0e, sid NN(0, since has normal columns. This
implies P [MO),,]] 21] < exp
1² 20
which in turn implies P maxi,j [M+), 2 S Eng P [M)), 2 1 S
53

---

Grokking Beyond the Euclidean Norm of Model Parameters
1.0
100
Nmin for recovery
80
0.8
coherence
0.8
80
60
Z
60
0.6
40
0.4
40
0.4
20
0.00
0.25
0.50
0.75
1.00
25
50
75
100
T
Sparsity level (s)
(a)
(b)
Figure 19: (a) Cohence (MM", 4) as function of E (0, 1) (b) Minimum number of samples for perfect recovery (relative recovery emorr < 10-9) for n = 102 as a function of the sparsity level E [n]] and coherence parameter E (0
nN exp
12
Using & = 2C2
(l((N) nt
with C2 2 1, we have t² = 2 (xddT) In a. with 1) (n.N)2-2C3 so
20
n. exp -0
12 17.
We also have the following theorem from Rauhut (2010) about the RIP of such matrix.
Theorem G.S. Let M E N be Gaussian Bernoulli random matrix. Let n, 6 (0.1) and assume N 2 C8-2 In (n/s) + In (1/n)) for universal constant C > 0. Then, (M) S 8 with probability at least n)
In the rest of this section,
control the incoherence, we generate M for given N by taking the first Nj = min([=NNJ,,) rows (with S r S 1. default to 0) from the first columns of & and the elements the remaining N2 = N Ni rows iid from (O 1/n)
+
I xn with MN.:: sid (000 The higher (and so N1), the less the MN
that X MΦ
MN
incoherence between the measures (columns of M and Φ. т correspond to a full random gaussian M. and correspond to the maximum incoherence, while correspond M, E (#:.))e()) for all i E (N), and correspond minimum incoherence (coherence of 1).
llo S s, and set b" $a. For a given s, we generate a random vector a" vid (0,1 such that ||a"|o
We generated & by performing QR decomposition random Gaussian matrix and taking the part.
For the noise & E RR,, we use ~
sid NO0, with SNR
ala" ll No
= 10 (signal to noise ratio).
G. 1.3. IMPACT OF COHERENCE ON GROKKING
Convex Optimization We fix 1 and solve for different (N,s77 s, 7) the problem (Pi) using the cvxpy library. As and/or increases, Nmin (8, T), the number of samples needs for perfect recovery increases (Figures 19 and 20). When converges to 1, ()) - for all s. The enror in those figures is the relative recovery error ||a a a²| [||///|a""]] This error is usually of the order of 10-6 giving us basis for comparison with other methods.
Gradient Descent Here, we also observe that the generalization time and the generalization delay increase with while the generalization error decreases with (Figures 21 and 22). For < n, when - 1, the generalization time 122 -> cc. This is because each measurement captures a single view (component) of b' = &a*. and this makes it impossible to find the optimal by solving the equation M = y" for N < n (by any method whatsoever). On the other hand, as
54

---

Grokking Beyond the Euclidean Norm of Model Parameters
10
3
10"
100
Error 0.0)
Sparsity level (s)
Sparsity level (s)
Nmin for recoveny
10
76
10 2
61
60
10
10
31
10-4
16
10-4
25
50
100
11
21
31
$1
61
71
81
91
101
110
20
40
60
80
100
Number of measurements (N)
Sparsity level (s)
Number of measurements (N)
109
100
Sparsity level (s)
Sparsity level (s)
Nmin for recovery
Error (T 0.2)
10-2
10
61
60
10
46
10
31
40
10
10
20
25
50
75
100
11
31
41
51
71
82
91
101
110
20
40
60
80
100
Number of measurements (N)
Sparsity level (s)
Number of measurements (N)
Error (T 0.5)
10a
Sparsity level (s)
g8
100
Mmi for recovery
10-2
Sparsity level
76
10-2
61
10"4
46
10
60
31
40
10
16
10-6
20
11
21
31
51
61
71
81
91
110
20
Number of measurements (N)
Sparsity level (s)
Number measurements (N)
10
89
109
100
Sparsity level (s)
Sparsity level (s)
min for recovery
Error (T 0.8)
10
76
10"2
80
ErToT
10-4
46
10
60
31
10
16
10-
0
25
50
75
100
31
51
71
91
101
110
20
40
60
80
100
Number of measurements (N)
Sparsity level (s)
Number of measurements (N)
86
100
Error (T 1.0)
10
Sparsity level (s)
Sparsity level (s)
10-1
Mmin for recovery
76
95
10-1
61
10-3
90
46
10-5
31
85
10-7
10
80
25
50
75
100
31
41
51
61
71
81
91
101
110
20
40
60
80
100
Number of measurements (N)
Sparsity level (s)
Number of measurements (N)
Figure 20: Relative error ||a - a" |//||a [|a"] || as a function of the number of measurements N, the sparsity level E [n] and and coherence parameter E (0, 1), for n 102
55

---

Grokking Beyond the Euclidean Norm of Model Parameters
- O. M becomes completely random, and every measurement captures a distinct "view" of b", giving the best possible generalization time for the data size considered. The error ||a(ta) - a ||la/|aa"]] at generalization (12) as function of
N and r has the same shape in the convex programming. We use (n, a, B) = (10?, 5, 10- 1005, 10-6) in the experiments.
G.2. Matrix Factorization
G.2.1. THE PROBLEM
Given a low rank matrix A' E R measurement matrix X R-Vxn1m;; aim to solve the following problem
for A E RMixaa; (P3) Minimize rank(A) subject to (eoe)A (X) y*1l2 S c; and more precisely, its convex relaxation
(P4) Minimize || = E; oi(() (A) subject to [[Fvec(A) (X) y 2l2 S E; where y = Fvec(AA)) (X) + E are the measures
and é an upper bound on the size of the error term & E RR |IE S.
Matrix Sensing
Matrix sensing seeks to recover a low rank matrix A E RRexxn2 from N measurement matrices {X; Ru X HE(N] and measures y (tr(X, sENNT We have y tr(x, A+) vec(X,)) vec(A*)
F Fvec(A A ()) (vec(X.)). This gives us compressed sensing problem, with the signal vector vec(A*) E Ruina and the measurement matrix X = (vec(X,))ece vec(A*)= = ① vec(E"); where veo(E*) E R is sparse since || veve Illo rank(A") S min(m n2) & 01712; and ① == V* @ U E Raidaxnina has orthonormal column since o = (V+ V*) & (U-TU) = Ianes
RNmmm In fact, under full SVD A' = U-E'V+T, we have a"
Matrix Completion For a matrix completion problem with matrix A E R we have N measurement vectors
X (1) X (2)
E R
x R and measures y X (I)T A-X(2)
X
x(") vec(A*) Ave (x{² x(")
X ( This gives us compressed sensing problem, with the signal vector vec A E R"1n2 and the measurement matrix X = ( ( E Nxnrne Standard matrix completion is usually defined as recovering missing elements of matrix from its incomplete observation. This is equivalent requiring x(() to be selection vectors for all & [K] (with K = -2,,ic.
i.. y = (2)
x vee(A*) = Fv))(A))
x(k) is the s(i, k)th vector of the canonical basis of R for certain s(i, k) E [mg]. This make each X; X (2) @ x(() a selection vector in R", and X ( x selection matrix in R so that y = A ((..)) Vi E [M]. So, in this formulation, each x (k) is a sample from the columns of Ing. Note
((()) yJ, that is y* = =y ( T less standard formulation of the matrix completion task requires each x!!)) be a sample from an orthonormal basis, i.e., x(())
that under a change of basis (() (()) x (k)), we have y; = (
is sample from the columns of V(k)_ R with (())T(( = In.. We let x be the s(i, 1-th column of (k) for certain sii, E [na]. Then y A ((..)))))))))))))))))))))))))))) with A A x1 v X2 ( So, any result state of A in the standard formulation where the measurement vectors are selection vectors is valid for the matrix A
Assume the target matrix A has rank r. Then it has r(n1 + n2 r) degree of freedom 10 and we need to observe at least r(n, n22 - r) entries for perfect recovery. This bound can be improved by considering the structure of A Let A = ~ V be the full SVD of A". As observed above, we are dealing with a compressed sensing problem with the signal vector a' = vec(A*) & vec(E"); where vec(E") E R=1n2 is sparse and &b V < U E [R"1ngg" X has orthonormal column.
G.2.2. THE CONTROL PARAMETERS
In this sub-section, we assume standard matrix completion. But the theories outlined here also apply to the general framework. The theory gives the minimal number of observations that guarantee A to be unique solution to problem (Pt) and allow perfect recovery of A' with fewer samples (Candès & Tao, 2010; Candes & Recht, 2012; Chen et al.,
2014). Generally, the lower bound on N has form N 2 C' max C2 (n log Ci (max n22) )) + log _ where 7)
1)
is the percentage of error (i.e N guaranteed perfect recovery with probability at least - 7), C1 > 0, C2 > 0, C3 > 0 are
1he first r columns of U form an orthonormal basis for a r-dimensional subspace of RT (the columns space of A Specifying this requires r(R1 - r) parameters. Similarly, the first columns of V form an orthonormal hasis for a r-dimensional subspace of R (the rows space of A ), and specifying this requires r(12 - r) parameters. The non-zero singular values are independent parameters. Thus, specifying them requires parameters.
56

---

Grokking Beyond the Euclidean Norm of Model Parameters
N=10 10
N=20
10a
1.0
10
10
58
-0a
10
27
10
0.
Error
||a 122 11x
|| y 112
fly "122
08
1|al
a
a|a 1l
24
10
||a
1| / ||a 11
0.
64
10
10
02
81
10
01
10
10
101
10
10a
104
103
10c
100
10
102
10a
10f
102
10c
Steps (t)
Steps (t)
N=30
N=40
10a
1.
10
100
09
23
10-1
0.7
Error
||
-y" 112
11y ll2
96
|
-y
ll2
1b
112
08
10-2
||a00
a 1l2
112
10
|
*12
||
24
04
10-1
10
91
0.1
10
10a
102
10
10a
10a
105
10"
10
107
10f
10a
10
Steps (t)
Steps
(t)
N=50
N=60
10a
29
48
0a
10-1
0.7
Error
||a th
y y'1lz
11y 11
28
|Ka tt
y
"lla
1y llaa
84
10
1|200
1l>
1b
10
||a et
||a
11
03
04
10
23
10-3
03
-01
10
101
102
103
104
2.
10-4
10
10a
10a
10f
10F
10c
0.0
Steps (t)
Steps
(t)
70
N=80
109
20
09
25
08
10
10
Error
11
11 112
96
06
10-2
H
10-2
03
||a
11
a
11
10
83
10-3
||a
-y
112
11
11a
03
02
||a0t
a
1112
||a
112
10-4
10a
101
10a
10a
10a
10
109
3.0
10-4
10"
103
102
103
10
10a
10
0.0
Steps (t)
Steps
(t)
N=90
N=100
10a
2.
10a
00
28
08
10-1
10-1
Error
56
06
10-2
25
10-2
05 |
10-1
||Xa 19
112
11y 112
10-3
||Xa
-y 112
lla
03
||a
a
11,
1b
52
||a
a
112
||a
l1
02
10
10
10a
16a
10a
10a
10a
10a
10"
10a
10a
10a
10a
10
10
10a
Steps
(t)
Steps
(t)
Figure 21: Training and recovery error as function of the number of samples N and the coherence parameter 7 E [0,
57

---

Grokking Beyond the Euclidean Norm of Model Parameters
10a
Ixxa y 1l 1I 1b
0.0
10"1
T- 0.3
0.4
T=08
T= 0.7
r= 1.0
20
60
80
N
10
T- 0.0
10
30
T- 0.3
1.0.0
T=0.1
10
10"
T-02
0.8
T=0.3
101
0.7
1.0
1-0.4
1-0.8
107
T- 1.0
1-0.7
10
10
1.0
1=0
Figure 22: Training and error Xa(()) y" 1 ll and recovery error (|a(ta) a a2112/ /a//|aall (along with t and t2, the memorization and the generalization step) as function of the number of sample N and the coherence parameter [0, 1].
constant, and C a universal constant. For example, in Candes Recht (2012). (C1, 1.2, 1) for small rank max(n1, n2) ).2 and C2 1.25 for any rank. The term max(n1, n2) ) log (max(n1: R2)) is due the coupon collector effect since to recover an unknown matrix, one needs at least one observation per row and one observation per column (Candes & Recht, 2012).
Definition G.6 (Random orthogonal model (Candes & Recht, 2012)). For given r, generate two orthonormal matrices U* E RR and V* E RraaXr with columns selected uniformly at random among all families of orthonormal vectors; and diagonal matrix Σ' with only the first diagonal element non-zero (with no assumptions about the singular values"! ! then set = -*****
We have the following result about the standard formulation for such matrices in the absence of noise.
Theorem G.7 (Theorem 1.1, Candes & Recht (2012)). Ler A E R be matrix of rank sampled from the random orthogonal model, and put max(n1, R2). Suppose random. Then there are numerical constants C and such that if N 2 Cn5/Ar log (n), the minimizer to the problem (Pa) is unique and equal to A with probability at least 1 1- c/nn;; that is to say, the semidenite program (P4) recovers all the entries of A with no error. In addition, if r S n 1/5 then the recovery is exact with probability at least - c/n²² provided that N 2 Cn 6/5, log (n).
x observe N entries of A with locations sampled uniformly at
Assume for example A (n1) (uz) for (k, E) E [n1] x [m2]. Even this matrix ranks at 1, it has only zeros everywhere except at position (k, f), e e₂ inputs. The only way to guarantee observation of the input position (k: {) is to choose measurements coherently with its (m1). This idea is formulated more generally below.
have very little chance of reconstructing it in high dimension by observing portion of its
singular basis ej @ ee("))
(na)
Definition G.8. Let U be a subspace of IR" of dimension and Py be the orthogonal projection onto U. Then, the coherence
of U vis-a-vis basis {u (n) 2012).
JsE[M) is defined by (uU) - max; ( We have < ((U) S n/r (Candes & Recht,
For matrix A T X under the compact SVD, the projection on the left singular value is x TU
I Unless otherwise specified, we default nonzero singular values 1 1/vF.
58

---

Grokking Beyond the Euclidean Norm of Model Parameters
and ||U x which considers each matrix entry.
| for all (similarly for the right singular value). We have the following definition of coherence,
Definition G.9 (Local coherence & Leverage score). Let A UEVT rank r. The local coherences of A are defined by
E R x be the compact SVD of matrix A of
(A) = 71 T
U
12
=
11
Vi E [m]]
r
(122)
(A) = 112
m2)
12
1/2
Vj E [m2]
r
r
12 are the leverage score of A (Chen et al., 2014), which indicate how "aligned" each row column of the original data matrix is with the principal components (the columns of U V). For each row i, (A) measures how much this row vector projects onto the subspace spanned by the first r left singular vectors in U. Rows with high leverage scores contribute more to the low-rank structure of A and are more "influential" in representing A. Similarly, V2 (A) measures the coherence of each column j in A with respect to the low-rank subspace formed by the right singular vectors in V. High values indicate columns well-aligned with the principal directions of A and play a significant role in capturing its structure. Matrices with uniformly low coherence scores have rows and columns that are evenly influential. contrast, matrices with high coherence scores for certain rows or columns have a few specific rows or columns that dominate the low-rank structure.
with 14, for row i and vj for row j. The quantities ||U c (m))) 1 and c
In the general formulation, this definition can be extended to the set from which the measures are chosen. But in general, it leads back to the standard formulation under the change of basis.
Definition G.10 (Generalize local coherence & Leverage score). We generalize the notion of coherence to any arbitrary set of vectors ( = {u [m) (u(")) ]E(N)) E (R"1XM, and ((m)) = =[v]-a) JENN E R2axNN, and defined the generalized local
coherences as
(A) = m (m))
Vie [Nl. (2(A) = n2 V y
Vj E [Na]
(123)
N2) = (21, u, u ((a) ((r)) J5t). We can write u fm))) = p(t)e!"))) and vi p(o)ef matrix from the canonical basis to U(mi) and (mz) respectively. So
Suppose the sets U(m) and ((a)) are be orthonormal basis (i.e.
(m2) (m2) = 8ü and
(ma) with p(t) e Ruexne the base change
m4(A) = 71
n1 J e 11 = mm(AA
[i
r
(124)
(A) = 12
(1z)
12
n2
e
(mz) m2)) = (A(A)
Vj [No]
X2 2(2) = P(P PO (2 Vv) = ÜΣVT. That said, any result stated in the standard formulation for A is valid for À under the general orthonormal formulation.
with À = A X1 (())
Candès & Tao (2010) and Candes & Recht (2012) used mainly an upper bound Ho on p and vs 140 2 max (max,g(n.) (A*))) max(e|p2) (A)), and define a constant in such that the max: ( V To =
maxij U.A.A
EUU.A V j.k S (u((A=))((A)) < Viinn Ho for all i, we can just take μn Z HOVE. From this, Candes & Recht (2012)
M1 ming' Since Σx U.AVV J.K S VIx U E V j.k
| la =
Viiing
show that if the coherence 1o is low, few samples are required to recover A".
Theorem G.11 (Theorem 1.3, Candes & Recht (2012)). Let A E R be a matrix of rank r sampled from the random orthogonal model, and put n max(ni, Suppose we observe N entries of A with locations sampled uniformly at
3 random. Then there are numerical constants C and e such that if N 2 C max (nt. Po Pl, Monk) arr3log (n) for some
3 > 2, the minimizer to the problem (P4) is unique and equal to A' with probability at least 1 c//².. In addition, if then the recovery is exact with probability at least 1 c/n² provided that N 2 8/5 (n).
r S n. /5 n5/5//p. /po.
Chen et al. (2014) show that sampling the element at position (i.j) with probability Pi = 2ips + v2) allows perfect recovery of A with fewer samples, and called such sampling strategies local coherence sampling.
59

---

Grokking Beyond the Euclidean Norm of Model Parameters
Theorem G.12 (Theorem 3.2 and Corollary 3.3, Chen al. (2014)). Let A [Rmx be a matrix of rank with local coherence (Pt, U/ Hi There universal constant co2 C1, C1,C2 02
such that each element (i,j) is independently (14+2;)r log" (nn +22) .1] min((s, m2)) j. then A is the unique optimal observed with probability Pij Z max min Co min( m2)
solution of the nuclear minimization problem (P4) with probability least 1 01 c/(n1 + 2 number of sample N O (max((),,1)) n22) log + n2)).
Given N and E [0, 1), to control the coherence,
For matrix factorization, we select the first N1 = -N examples with the highest values (u;(A') +1 (A"), and select the remaining (1 >> examples uniformly among the rest. The positions selected are one-hot encoded in dimensions Bi (for row positions) and n2 (for column positions) have x(1) and (2), respectively.
For matrix sensing, we generate (1) (resp.
x by taking the first Ni min( [rM],,n (resp. N1 min[[TN], n2)) rows from the first columns of U* (resp. V*) and the elements of the remaining N₂ = N Ni rows iid from the Gaussian distribution (0,, o²) with o = 1/nz (resp. o 1/n2).
The higher 7 (and so N1), the less incoherence between the measures (rows of X x X and $ V U.
G.2.3. IMPACT OF COHERENCE ON GROKKING
Linear programming We fix n1 = nt2 = 102 and E = (no noise) and solve for different (N,r,7 7) the matrix factorization problem presented above using standard linear programming (we use the cvxpy library). As r and/or T increases, the number of samples needed for perfect recovery decreases. The relative recovery error ||A A ||////|| A obtained is usually of the order of 10-6 and gives us basis for comparison with other methods. We do not include figures save space.
Gradient Descent Unlike compressed sensing where large values of r detrimental to generalization, here, performance improves, and the number of examples required to generalize decreases exponentially, does the time it takes the models to do so. See Figure 23 for (n1, 12, a, B, {) (10, 10, 2, 10 10 10 1009). 9). Note that here, for matrix completion, for fixed т, we select the first rN examples with the highest values of Hu( (A(A*) + remaining (1 - N examples at random, uniformly.
and select the
N=70
N=70
1.0
0.9
10
0.8
10
0.8
0.7
0.7
Error
||xee(A (1)) -y"ll2 Ily ll2
0.6
IIXvec (0)
y ll2 By ll2
10-3
(14(0) - ||F
||A I|F
10
1|A(H
ll
IIA
I|+
0.4
10-5
10-5
0.2
0.1
0.1
0.0
0.0
10a
10a
10a
107
1
2
4
Steps (t)
Steps (t)
Figure 23: Training error |X vec (t) of sample N and the coherence parameter E (0,
-y"122/ Ily ||2 and recovery error [|A A A ||e a function of the number
H. Additional Experiments
H.1. Sparse Recovery
Optimization landscape We look the landscape of the solution in the context of sparse recovery. Let I = {i
[n] a * 0} the support the ground truth signal a"; u(t) = ||a{") ll2 and v(t) = |a |a*))r/le llla lla be the norms (4) restraint on its indexes in and outside T, respectively. Figure 24 shows how (*) first converges to the least square
60

---

Grokking Beyond the Euclidean Norm of Model Parameters
solution (memorization), and from the least square solution a' (N large enough) or a suboptimal solution (N too small). After memorization, when N is large enough, u(t) converges because the components of (() that are not in / are shrunk at each training step until they all reach 0 (Figure 25). This convergence is impossible if the & regularization strength is 0 (even if & is used). For the experiments of this section we use N E (20, 30, 40, 50, 60, 70), for (n, 8) = (100, 5) and (a,, (3) = ,0--9)
while u(t) converges the norm of a*. This is
100
0.06
0.05
60
v(e
0.04
50
(least square)
(optimum)
tg
20
0.01
t> = 95
437
70
0.00
t. 999000
10
0.00
0.02
0.04
0.06
0.08
0.10
0.12
0.14
0.16
u(t)
From initialization to least square solution à (memorization), and from least square solution to at (N large enough) or a suboptimal solution (N too small). The steps t and i/ are different from those introduced above to measure memorization and generalization (respectively). They are just means of tracing the evolution of training here.
Figure 24:
Scaling the Learning Rate a and the Regularization Strength We solve the sparse recovery problem using the subgradient descent method with (n,, 8, () (10², 5, 30, 10-6) for different values of a and B. As expected, larger and/or 3 lead to fast convergence and do at suboptimal value of the test error (Figure 26). However, small values require longer training time to plateau and generally do so at lower value of recovery error (grokking). These are the experiments used to produce Figure 4.
Scaling the Data Size N and the Sparsity Parameter We solve the sparse recovery problem using the subgradient descent method with (n, S. a. 3) = (102, 10-6, 101111000000 for s E {1,5, 10, 15) and N {10, ...... 100). See Figures 27 and 28 for the results. Smaller requires a smaller N for generalization.
H.2. Matrix Factorization
We optimize the noiseless matrix completion problem using the subgradient descent method with (n1, n2, N. (10, 2, 70, 10-6) for different values of a and 3. As expected, larger and/or 3 lead to fast convergence and do so at a suboptimal value of the test error (Figure 29). These are the experiments used to produce Figure 6.
H.3. General Setting
In this section we optimize functions of the form f(0) 9(0 + 30),,, where g is the square loss or cross-entropy loss function of the considered model on the training data, 0 the set of model parameters, and h regularizer applied to 6. can be the standard fp norm or quasi-norm of f, the sum of the nuclear norms of each matrix in 0 (i.e. f.), etc. By normal
initialization for parameter A E R XR2, we mean A (1) idd (0, 1 1/n1). For the experiments of this section only, we used
61

---

Grokking Beyond the Euclidean Norm of Model Parameters
0
106
0.10
103
0.08
Steps (t)
0..66
102
0.04
10
0.002
0.00
100
11
16
21
26
31
36
41
46
51
56
61
66
71
76
81
86
91
96
100
Dimensions (n)
Figure 25: Convergence of a, to a for each i [n). Here (n, 8, N) (100, 30) and (a, (10-1, (10-1,10-5),
(t)
0=
0.005
0.01
107
20-1
107
10-
10
Error
10
10
10-
10
10-1
||Aa
-y
11
Iy
11
10
|Ka
11
11
1b
||a
1:=
1|8
1
12
||a
lla
10-3
10-1
10:
10a
10a
10"
10
10a
10
10'
0.0275
a=
10a
10a
10-1
10"²
Error
30-2
20-
10-2
[Ke 11
11 la
20-1
10-4
10-1
1a"
12
lallla
20
11021
10-4
0.0825
0.1
109
107
10-4
Errr
10
||Ka
112
11 112
20-
10-2
20
[|a
11:
10-1
10
1b
[8
10:
10a
3p-1
10:
10a
10a
107
Steps (0)
Steps (t)
Figure 26: Training and recovery error as a function of the learning rate a and the erreglaarinztionn coefficient 3.
62

---

Grokking Beyond the Euclidean Norm of Model Parameters
s=1
s= 5
10
100
10
100
10-1
10
70
70
Error
||a in
y ll2 / Yyy ll2
Error
|| y 1l2 / Bly ll2
||
ll2
||a
1|2
50
10
||
1|2
|aa
llz
50
10
40
40
10-3
20
20
10-3
10
10
100
10a
10a
10a
10
105
10c
100
10a
10
10a
10a
10a
10
Steps (t)
Steps (t)
s=10
s= 15
10a
100
10
100
10-1
10
70
70
Error
Xa hn
y ll / 1yy 1l2
Error
Ia (2)
ll2 / 1l ll2
10-2
10
||a00
112
||a
1l2
||ath
112
||a
ll2
10-3 -
40
10-3
40
10
20
10
20
10
100
10a
102
10a
10
10
10
100
10a
10
10
10
10
10
Steps (t)
Steps (t)
Figure 27: Training and recovery error as a function the sparsity level 8 and the number measurements N.
/ [10 112
10'22
[Ktal V
13 1/3
10"7
60
1.00
100
1.7
127
47
1.7
1.5
100
300
Figure 28: Training error [Xa(t)) -yy 12/ |ly" [|2 memorization, recovery error [|a(t)) - aa|l/\aaa* at generalization,
(1)
memorization step I, (smaller / such that [Xa(t) y" 10-1), and generalization step (smaller & such that (t) a000>aalll///// a*||a/||a" 22 < 10 or the maximum training step).
63

---

Grokking Beyond the Euclidean Norm of Model Parameters
a= 0.0825
a = 0.1
0 0.75
10
10
10
Error
10
[/ y
10""
[[xvecIA -y -y'lls
10-3
IKvecIA " - 11
1 AA'
||r
IM II
111 -1 1rr
30
11 1 11
IM
10-4
10-3
10
10a
10a
10a
107
10a
107
10a
107
10
102
102
10'
0.825
1.0
1.5
30
211
Error
10--
!/ " -Y 19
10-4
[[VeCIA -y 112
10-3
IKvecIA -y y 11a
324
1 AA [|r IN
-1 -1 11r 1A 1rr
11th A 1rr 111
10-5
10-0
107
10a
10
162
10
10:
162
10a
Steps
Steps (t)
Steps
Figure 29: Training and recovery error as a function of the learning rate a and the l, ee-reuulaiization coefficient J.
Adam as the optimizer, with its default parameters (as specified in PyTorch), except for the learning rate.
1.3.1. ALGORITHIICC DATASET
We consider addition modulo p = 97 described Section 2.1 with Train 40%. For MLP, & and l. have the same effect
on grokking f2, i.e., large a& values lead to faster grokking. See Figure 2 for E,, and 30 for & and lz. Note that for the MLP, the logits given by yxx) (( + w σ ((() + w(1) (E E(zz))), with ((2) max(z, 0). E E Rexd1 ((I) Rdaxd,, b(1) E Rd,, (2) E Roxds, and b (2) R", where the embedding dimension. We use d1 d2 20.
0.001
a 0.01
a = 0.1
1.0
1.0
1.0
Accuracy
10-5
10
0.6
Train
0.6
0.6
0.4
Test
10-
0.4
10
0.4
10
0.2
0.2
Train
0.2
Train
10m9
Test
10
Test
10
2500
10000
5000
7500
10000
2500
5000
10000
Steps (t)
Steps (t)
Steps (t)
a= 0.001
a 0.01
a 0.1
1.0
Accuracy
0.8
10-1
0.8
10-
0.8
10-1
0.6
Train
0.6
Train
0.6
Test
Test
10
0.2
0.2
0.2
Train
Test
0.0
10-7
0.0
10-7
0.0
10-
2500
5000
7500
10000
0
2500
5000
7500
10000
2500
5000
7500
10000
Steps (t)
Steps (t)
Steps (t)
Figure 30: Training and test accuracy of MLP trained modular addition with & (top) and & (bottom) regularization for different values of the learning rate and the & coefficient 3.
64

---

Grokking Beyond the Euclidean Norm of Model Parameters
H.3.2. NON LINEAR TEACHER-STUDENT
We consider teacher yx) (²x((Ax)) from R Rf with hidden neurons (A* E RRxx and B [Rex)); where (()) max(2, 0) and x, rB" isd (00, 1). We i.idd sample N inputs output pair Drain {(x, y" (x:))) and optimize the parameters 0 (A, B) of student yx(x) = Bø( (xx on them, starting from normal initialization, with the loss function g0) 2N Σ Ilye (x,) (xx))11 and different regularizer h(0), lp for (1,2,*) *). For all these regularizers, the smaller ag is, the longer the delay between memorization and generalization. See Figures 31, 32 and 33 for
an experiment with (d, r, N) (100, 500, 2, 10a).
0.0001
0.00055
0.001
10
10
12.5
Train
12.5
Train
12.5
Train
Test
Test
Test
10-1
10-3
10-2
Loss
7.5
7.5
7.5
10-4 CD
10 COL
10- 00.
5.0
5.0
5.0
2.5
2.5
2.5
10
10
0.0
10-2
0.0
10
0.0
10-7
10a
10
10a
10a
10-
10a
a
0.0055
a' 0.01
0.055
60
2000
10
12.5
Train
Train Test
Train
10
30
10
Test
10.0
1500
40
10
Loss
7.5
1000
10-
10-9
10-
500
2.5
0.0
10-
0
10
0
10
10
10a
10a
10a
10a
102
10
10a
10"
Steps (t)
Steps (t)
Steps (t)
Figure 31: Training and test error two layers ReLU teacher-student with & regularization, for different values the learning rate a and the coefficient
H.3.3. DOMAIN SPECIFIC REGULARIZATION
Neural Networks (Raissi et al., 2019) leverage prior knowledge from differential equations by incorporating their residuals into the loss function, ensuring that solutions remain consistent with physical laws. Sobolev training (Czarnecki al., 2017) generalizes this idea by incorporating not only input-output pairs but also derivatives of the target
function. More precisely, given input-output pairs {(x,, ( (x: (x:))e(N) along with known derivatives
(x)
X=X;
iE[N]
for k E [K]. the goal to train a neural network (x(x) that approximates both the output and its derivatives. The loss function extends the standard mean squared error (MSE) include Sobolev penalties:
N vo((x) - yx::))
akyo Oxck
g
((0) =
(x:)
Fx (xi)
(125)
N
N
1=1
i=1
data
Sobolev penalty
The hyperparameter 3 controls the contribution of the derivative alignment term. This penalty ensures that the model not only fits the data but also respects known smoothness constraints or differential structure, which is crucial in physics- based applications (Lu et al., 2021). We consider the two layers feed forward teacher of Section H.3.2, and optimize the parameters = (A, B) of a student using the sobolev objectify for K = 1, Qy' (x) = B diag (d' (d((*))) A". The smaller is aB, the longer is the delay between memorization and generalization. See Figure 34 for an experiment with
âx
(d, r, c, N) (100, 500, 2, 102).
65

---

Grokking Beyond the Euclidean Norm of Model Parameters
a 0.0001
0.00055
a 0.001
18-1
1811
1014
15
Train
15
Train
15
Train
10-a
30-1
10
Test
Test
Loss
10
10
10
CO
CO
00
10-9
10-
10-
0
10-1
10
2017
10a
10a
10
101
10a
10
10
10a
10
a
0.0055
a 0.01
a 0.055
50
15
Train Test
Train
Train
10a4
40
Test
10a4
1500
Test
101h
Loss
10
10-3
30
10-3
10
1000
20
500
10
10
10
10
101
10
107
101
107
105
10
107
10a
Steps (t)
Steps (t)
Steps (t)
Figure 32: Training and test error two layers ReLU teacher-student with & regularization, for different values of the learning rate a and the l2 coefficient B.
a = 0.0001
a = 0.00055
a = 0.001
12.5
Train
12.5
Train
12.5
Train
10
10
10-4
10.0
Test
10.0
Test
10.0
Test
LOss
7.5
7.5
10
10t1
10
2.5
2.5
2.5
10
10
0.0
0.0
0.0
10a
10a
10a
10a
10a
10a
10
10a
10a
a
0.0055
a 0.01
0.055
10a
10aa
10-4
12.5
Train
Train
Train
1012
1500
10.0
Test
Test
Test
10
10
20
Loss
7.5
30
1000
10
10
10
5.0
20
500
2.5
10
10t4
10-4
10-1
10-1
105
103
10a
10a
10a
Steps (t)
Steps (t)
Steps (t)
Figure 33: Training and test error two layers ReLU teacher-student with E, regularization, for different values of the learning rate a and the f, coefficient B.
66

---

Grokking Beyond the Euclidean Norm of Model Parameters
0.0001
0.00055
0.001
18
20
Train
Train
Train
Test
Test
Test
10
10
10
10
Loss
10
CO
10 CO.
10 CO
10-
10-1
10-1
10
1011
10
10a
10
10a
10a
10a
10t
10a
10a
a 0.0055
a 0.01
a 0.055
2000
Train
Train
Train
1014
10
20
Test
Test
Test
10
10-1
40
30-
1500
10
Loss
10a4
20
10a4
10
500
10-4
10-4
10
10a
10a
10
10¹
10a
10
10a
10a
Steps (t)
Steps (t)
Steps (t)
Figure 34: Training and test error two layers ReLU teacher-student with Sobolev training, for different values the learning rate a and the l coefficient B.
H.3.4. IMAGE CLASSIFICATION
We optimize the parameters 0 = (A, B) of a model yo(x) = Bø (Ax) on = 1000 samples the MNIST dataset. Figure
35 show the results for 4 the result for & and f, are similar.
0.001
0.005
0,01
1.00
10
1.00
10
1.00
Train
Train
Train
Accuracy
0.75
Test
30
0.75
Test
10
0.75
Test
10
10
10a
10
0.50
0.25
0.25
0.25
10
0.00
10a
10
10a
104
10
10a
10a
10a
10
10a
10a
102
10a
10
Steps (t)
Steps (t)
Steps (t)
Figure 35: Training and test accuracy of a MLP trained on MNIST with & regularization for different values of the learning rate & and the l coefficient, |
67


Reviews:
[{"review_id": "ggXrFER3JU", "reviewer": "Reviewer_WymT", "summary": "The authors study delayed generalization (grokking) in sparse/low rank recovery tasks. The authors focus on the transition from an overfitting solution to a generalizing solution, both in linear sparse recovery and in low rank matrix factorization. For the linear cases, the authors derive the scaling of the grokking time difference. They show that for sparse tasks, grokking can be tracked and controlled by the $L_1$ or the trace norm of the weights or dictionary vectors.", "questions": "", "limitations": "", "rating": 2, "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 time and insightful comments. We appreciate the effort taken to evaluate our work and provide constructive feedback. Below, we address each of the points raised.\n\n**Weaknesses**\n\nWe understand that the paper is poorly organized. That is why, after submission, we completely rewrote it as follows: \n* All the proofs have been simplified, and we've included the proof sketch for each main theorem. For example, here's the main idea behind Theorem 2.1. We first show (using standard results from compressed sensing literature) that, given a matrix $\\tilde{\\mathbf{X}} \\in \\mathbb{R}^{N \\times n}$ and a vector $\\mathbf{b}^* \\in \\mathbb{R}^{n}$ of sparsity $s$ ($\\|\\mathbf{b}^*\\|_0 \\le s$), we can write, for any vector $\\mathbf{b} \\in \\mathbb{R}^{n}$, $\\|\\mathbf{b} - \\mathbf{b}^*\\|_2 \\le C_1 \\| \\tilde{\\mathbf{X}} \\left(\\mathbf{b} - \\mathbf{b}^*\\right)\\|_2 + C_2 | \\|\\mathbf{b}\\|_1 - \\|\\mathbf{b}^*\\|_1|$ where $C_1$ and $C_2$ are constants that depend only on $s$ and $\\tilde{\\mathbf{X}}$. After a short amortization phase ($t \\le t_1$), the term $\\| \\tilde{\\mathbf{X}} \\left(\\mathbf{b}^{(t)} - \\mathbf{b}^*\\right)\\|_2$ vanish (or becomes proportional to the noise if there is any), since we have $ \\tilde{\\mathbf{X}} \\mathbf{b}^{(t)} \\approx \\mathbf{y}^* = \\tilde{\\mathbf{X}} \\mathbf{b}^* + \\xi$. Then comes a second phase, during which $\\|\\mathbf{b}^{(t)}\\|_1$ converges to $\\|\\mathbf{b}^*\\|_1$ at $t_2$: our contribution is to show that $t_2 - t_1$ is inversely proportional to $\\alpha$ (the learning step) and $\\beta_1$ (the regularization strength).  \n* We made the body of the paper self-contained insofar as any contribution mentioned in the abstract and introduction is proven in the body of the paper theoretically and/or empirically (grokking time, effect of data selection on grokking, limits of $\\ell_2$, and grokking without understanding, ...). So, the results in the appendix this time are just extensions of those in the main text.\n* With the rewrite, the appendix has also been reduced from $\\mathcal{O}(70)$ to less than $50$ pages.\n* We have left the part on matrix factorization in the main text since the mechanics behind the generalization phase, in this case, differs from sparse recovery. In sparse recovery problems, we take into account only the iterate $\\mathbf{b}^{(t)}$ (especially its $\\ell_1$ norm), whereas, in matrix factorization, we take into account not only the singular value matrix $\\Sigma^{(t)}$ (especially its $\\ell_1$ norm) of the iterate $\\mathbf{A}^{(t)}$, but also its singular vectors $\\mathbf{U}^{(t)}$ (left) and $\\mathbf{V}^{(t)}$ (right), which makes the dynamic in the second case completely different from the first case. Also, in a matrix factorization problem, replacing $\\ell_*$ (the nuclear norm regularization) by $\\ell_{*/2}$ does not induce grokking (at least in the shallow case used in the theory).\n\nThanks for the suggestion about figures and parameters of interest. We will define a parameter $\\tilde{\\alpha} = \\alpha \\beta$, which is the main parameter in the theory ($\\alpha$ is the learning rate and $\\beta$ is the regularization).\n\n**What is the overlap between $\\mathbf{b}^{(t_1)}$ and $\\mathbf{b}^{(t_2)}$? What changes more, the norm or the direction?** Once memorization occurs ($\\mathbf{X}\\mathbf{b}^{(t_1)}=\\mathbf{X}\\mathbf{b}^*+ \\xi$), the direction of $\\mathbf{b}^{(t)}$ that explains the data does not need to change drastically.  The main change from $t_1$ to $t_2$ is the magnitude of $\\mathbf{b}^{(t)}$ in $\\ell_1$—i.e., it transitions from something closer to the least‐square solution’s norm $\\|\\hat{\\mathbf{b}}\\|_1$ down to the teacher’s norm $\\|\\mathbf{b}^*\\|_1$ (while still maintaining some form of memorization). So, $\\mathbf{b}^{(t_1)}$ and $\\mathbf{b}^{(t_2)}$ are closely aligned in direction; the main difference is that the $\\ell_1$ norm shrinks from $\\|\\hat{\\mathbf{b}}\\|_1$ to $\\|\\mathbf{b}^*\\|_1$. Hence, the norm changes more than the direction.\n\n**Can you comment on the fact that the test loss is always monotonic in your settings?**\n\nThe monotonic decrease in test loss is due to how we choose the hyperparameters, mainly the learning rate and the regularization strengths. In the case of sparse recovery, for example, the test loss is controlled mainly by $\\|\\textbf{b}^{(t)}\\|_1$, even more so after memorization, which takes only a few steps. After this memorization, $\\|\\textbf{b}^{(t)}\\|_1$ converge very slowly towards $\\|\\textbf{b}^*\\|_1$. When $\\alpha$ and/or $\\beta_1$ are larger than the values used in the paper, there are indeed a lot of oscillations in test loss. The same reasoning applies to experiments outside the linear scope (e.g., figure 5). We chose the learning rate and regularization strength to illustrate grokking. Larger values show oscillations (e.g., the Slingshot Mechanism [1]).\n\n[1] The Slingshot Mechanism: An Empirical Study of Adaptive Optimizers and the Grokking Phenomenon (https://arxiv.org/abs/2206.04817)", "claims_and_evidence": "The claims are both proven and supported by empirical evidence.", "methods_and_evaluation_criteria": "The problem is well set up and the evaluation criteria make sense fo this problem.", "theoretical_claims": "I only checked the main theorem 2.1, while the rest are somewhat demonstrated directly by the experiments.\n\nI would state that the paper contains 81 pages, and it would be unreasonable to check in its entirety in the allotted time for review.", "experimental_designs_or_analyses": "The main experiments are simple and done on convex problems, and are therefore reproducible with the parameters given in the paper.", "relation_to_broader_scientific_literature": "The contributions presented in this work are related to the broader scientific literature, in particular they relate to grokking in linear estimators, where delayed generalization can occur even in convex problems depending on certain parameters of the setup. They extend these ideas to more complicated settings where a clearly defined overfitting and generalizing solutions exist.", "essential_references_not_discussed": "No essential references clearly come to mind for the specific context of this paper, certainly not all grokking papers are cited but since the focus is not on feature learning but just on the transition between possible solutions. Perhaps the authors could contrast their work more with the feature learning works such as [1,2].\n\n\nRefs:\n\n[1] - GROKKING AS A FIRST ORDER PHASE TRANSITION IN\nTWO LAYER NETWORKS, Rubin et al., ICLR 2024\n\n[2] - GROKKING AS THE TRANSITION FROM LAZY TO RICH\nTRAINING DYNAMICS, Kumar et al., ICLR 2024", "other_strengths_and_weaknesses": "**Strengths:**\n\n1) The problem setup and topic are very interesting in my opinion, and are of interest to the DL community.  I have not seen works that focus on the relation between grokking and sparsity in this way, and I believe it can be a key contributor to this phenomenon in many real world settings.\n\n2) The results seem robust and correct, and provide a new way to look at progress measures of grokking, possibly in problems where the problem is not tractable analytically.\n\n**Weaknesses:**\n\nThe main and overwhelming weakness of this paper is its extreme breadth and length. The authors attempt to cover too many setups and end up compromising quality for quantity. By attempting to tackle sparse recovery, deep sparse recovery, data selection, low rank recovery, matrix completion, and extending to nonlinear cases, the main text cannot provide more than a shallow discussion for many of these cases, which leads to the following problems:\n\n- The paper is poorly written, both in terms of the number of typos/grammatically incorrect sentences.\n- When theorems are given, there are no sketches of the proofs, which should provide the intuition and deep insight as to why these results work and make sense, but instead just given in terms of full proofs in the appendices, which are $O(70)$ pages. This is not the correct format for a conference paper.\n- Many results and figures do not appear in the main text, but are referred to the appendices. This would be fine if not for comments such as *\"While grokking has been studied extensively, the impact of data selection on grokking remains largely unexplored, making this one of the first works to address this critical aspect.\"*  Either a topic is important enough to appear fully in the main text and discussed seriously, or it is not truly critical, it cannot be both. Results in appendices should be either trivial extensions that do not add any deep insights, or indeed proofs, when the sketch is given in the main text. Another example is the Deep Sparse Recover paragraph, which directly points to the appendix, as well as the Realistic Signals paragraph. All of these should either be removed from this work or discussed seriously in the main text.\n - There is no mention of the second effect that typically arises in grokking which is the non-monotonicity of the test loss while the training loss monotonically decreases. This is a minor point but is never discussed.\n\n\nThis weakness is the reason for my inclination to reject this paper at the current version.", "other_comments_or_suggestions": "**Partial list of typos and grammatical errors:**\n\n- The paper has no page numbers at the bottom, not sure how this happened.\n- L20 (right) - \"we always need an ℓ2 regularization term.\"\n- L23 (right) - \"the dynamics of\"\n- L37 (right) - \"previous art\"\n- L62 - “with an mlp”\n- L125 - \"number of samples needs\" should be \"needed\"\n- L163 - \"since for problem of interest\", should be \"for the...\"\n- L171 - \"If β2 is choose such..\"\n- L216 - \"number of measures N...\", should be \"samples/measurements\"\n\n- $\\tau$ is not defined before it is used in the text, and it is unclear if the error term $\\xi$ is always meant to be a constant vector or a random vector which is different per sample, etc.\n\nThere are more which I have not included here.\n\n**Suggestions:**\n\nMy main suggestion given the weaknesses and strengths of the paper is simple: divide this paper into at least 2 different works, one which simply focuses on sparse recovery/deep sparse recovery, and explains the related results in depth in the main text, including sketches of proofs, intuition, and dependence on the relevant parameters. Then a separate work that focuses on low rank estimation, or putting the low rank estimation in the appendix of the first work, since it doesn't seem to provide a much deeper insight except for replacement of the lasso norm with the nuclear norm. Then I would suggest a concrete rewriting of the paper to ensure that the grammatical mistakes do not repeat and improving the structure. I would also suggest that the authors consider describing the problem in terms of its invariant quantities that actually dictate the scaling, for instance $\\alpha \\beta_1$ is one of the main scaling quantities, not $\\alpha,\\beta_1$ separately. It would be smart to simply define products and ratios of the original parameters and study the problem in terms of these quantities, to avoid cumbersome notations and clearer takeaway messages.", "questions_for_authors": "**Questions:**\n\n1) What is the overlap between $b(t_1)$ and $b(t_2)$? What changes more, the norm or the direction? \n2) Can you comment on the fact that the test loss is always monotonic in your settings? in general it is possible to have a non monotonic test loss even for a convex problem if the transition from overfitting to generalizing solution involves a big change in the parameters (either alignment or norm), so I'm curious if you can explain why it doesn't happen in your case. In some figures (fig 5 for example) you do have it, but only when it is accompanied by a previous non-monotonic effect in the training loss, which is non generic.", "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": "lTkC4wjxZG", "reviewer": "Reviewer_8Ky8", "summary": "The paper mainly focuses on the theoretical approach to debunk the necessity of L2 norm for exhibiting grokking phenomenon. As an alternative, the authors suggest that sparsity of the solution space can be alternative condition to have grokking phenomenon in deeper layer (proven practically). The paper includes rich theoretical justification on a linear matrix fitting problem and demonstrates with 2-layer to 3-layer MLP with ReLU activation, PINN, and LSTM to justify their claim.", "questions": "", "limitations": "", "rating": 2, "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 appreciate the reviewer’s time and insightful comments. Please find our responses to the points below.\n\n**Claims And Evidence**\n\n* We worked in a setting where the notion of sparsity (resp. low rank) is well defined, $\\mathbb{R}^n$ (resp. $\\mathbb{R}^{n_1 \\times n_2}$). This is the number of non-zero elements of a vector (resp. the number of non-zero singular values of a matrix). If we decide to work, say, in $\\mathbb{Z}/p\\mathbb{Z}$ for some prime number $p$ (like in the first problem grokking was first reported on, modular arithmetic), then these definitions are not trivially valid, and the theory we developed no longer applies (indeed, the theory of sparse recovery and matrix factorization over finite fields is very complex and uses completely different tools from those used in $\\mathbb{C}$). This is why we can not answer the question, “Is the modular arithmetic problem to recover the removed half of the binary relationships sparse?” using our current theory and leave it as a future direction (see section F of the Appendix). Although for tasks such as the sparse parity task (on which grokking is observed [1, 2]) one can construct sparse MLP that fits the data [1], we are not aware of any work that has done this for modular arithmetic. [3] constructs a two-layer MLP with square activations that does modular addition, but the model weights constructed are dense. \n\n* The most crucial point of our work is to show that not only does the grokking mechanism go beyond the $\\ell_2$ norm, but also why it is necessary to do so. Indeed, we show that the time to generalize is inversely proportional to the regularization strength used and that the generalization error after convergence decreases with it. Moreover, we show that grokking can also be highly accelerated by simply choosing appropriate data samples. To the best of our knowledge, no previous work on grokking has done all these.\n    \n    Although our theory relies on simple frameworks that allow us to extract the laws justifying our results, we also show this empirically for different types of models and datasets, including the basic setup on which grokking was first observed, the algorithmic dataset. We acknowledge that developing a theory in a simplified framework and then testing it experimentally in a more complicated setup has its flaws. The point of our paper is not to say that the mechanisms behind linear models extend to non-linear models but to illustrate that grokking is nevertheless observed if we use the insights obtained theoretically on linear models. We'll try to make this point clearer in the final version of the paper.\n\n* We acknowledge that grokking can be caused by a number of factors. This is the point of our work, even if we focus on regularization. Other work (Grokfast, etc) focuses on other factors to amplify/reduce the Grokking delay, and we will discuss this in detail in the related works section.\n\n**Methods And Evaluation Criteria, Theoretical Claims**\n\nThe point of our work is not that our theoretical results on linear models **trivially** extend to non-linear models. We simply show that the mechanism behind grokking (and regularization) goes beyond the $\\ell_2$ norm: we show it theoretically and empirically in specific settings and just empirically in others (and leave the theory in such settings as future work).\n\n**Essential References Not Discussed** \n\nThanks for the references, we will discuss this in the related work section. These works differ from ours in that the first focuses on double-descent (which is not the point of our work) and the other on other strategies to accelerate grokking (we focus on different types of regularization and data selection).\n\n**Questions For Authors**\n\nWe excluded Transformer results on modular arithmetic as the paper was already too long. On Transformer, $\\ell_1$ and $\\ell_*$ still affect the delay between memorization and generalization.\n\n**References**\n\n[1] Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit (https://arxiv.org/abs/2207.08799)\n\n[2] A Tale of Two Circuits: Grokking as Competition of Sparse and Dense Subnetworks (https://arxiv.org/abs/2303.11873)\n\n[3] Grokking modular arithmetic (https://arxiv.org/abs/2301.02679)", "claims_and_evidence": "The paper has laid some theoretical works to disentangle the necessity of L2-type regularization and the phenomenon of grokking. Moreover, this leads to that L2 weight norm may not be a good indicator for grokking progression measure. I find potential significant issues in the manuscript:\n\n- Although suggesting that the “solution sparsity” behaves as alternative potential source of grokking is interesting and meaningful, there should be more detailed definition with rich examples what solution sparsity is. For example, is the modular arithmetic problem to recover the removed half of the binary relationships “sparse” (this is the first problem that grokking is reported on)?\n- There is a clear logic gap between the theoretical justification and the actual claim. I agree that the amount of theoretical works of the authors are respectable; the theory seems to be rich enough, but for its own. The theory deals with the problem with deep linear layers which is not commonly used in practice nor it is dealt in the typical literature on grokking. The way the claim that made on the linear layer connects to nonlinear experiment seems a large logical gap to me. And to address this, I would like to suggest a theoretical work on nonlinear layers.\n- It is already known that other factors than L2 norm can reduce grokking. For example, papers like Grokfast (arXiv 2024) have shown that enhancing a momentum-like term of the gradient descent orthogonally reduce the grokking delay. Therefore, the “connection between L2 norm and grokking” may not be as significant as the authors has initially claimed.\n\nUnless these issues are addressed in the rebuttal, I am afraid I cannot agree with the significance of the author’s claim.", "methods_and_evaluation_criteria": "The paper tries to justify its claim theoretically especially on matrix factorization and least squares problem, and then tries to relate its conclusions on the deeper layer neural networks. Although the grokking analysis on “deep linear problems” is interesting enough, I do not agree that the conclusion for this problem trivially relates to the nonlinear network case. I have checked multiple times throughout the theoretical section, but I have not found a theory that focuses on at least two-layer nonlinear networks. There is a missing gap between the theory and the evaluation and I believe that these two phenomenon can still be regarded as separate phenomena unless a proper discussion encloses the gap.", "theoretical_claims": "I have tried my best to check the correctness of the proof, but I must say I have missed many theorems especially those appear in the supplementary material. Although I agree that the presented mathematical justification indeed supports the author’s claim that Euclidean norm is not necessary for grokking, I do not believe this can be trivially extended to more complicated architecture, e.g., Transformers, where the grokking phenomenon is first observed. Moreover, I am still not convinced that the theory on linear problem is extendable to nonlinear experiments.", "experimental_designs_or_analyses": "If the Transformer-based experiments (the one like in the first report on Grokking) are presented along with the other experiments, it would be more nicely looking. This is not a cumbersome experiment.", "relation_to_broader_scientific_literature": "Deep linear layers are in fact has a domain specific usage, e.g., in kernel estimation problem (e.g., in KernelGAN) and other problems that require redundancy in parameter space for smoother training. However, it is not well known that grokking is significant in these model domains. There this paper might get its significance.", "essential_references_not_discussed": "- The discussion between grokking and double descent has been quite significant, and the authors might get better insight in the joint works, e.g., \"Unifying Grokking and Double Descent\"\n- Groktransfer (ICLR 2025) and Grokfast (arXiv) both addresses different acceleration strategies other than exploiting sparsity and l2 norm.\n- Papers like “Grokked Transformers are Implicit Reasoners” discusses large scale experiments, too. I believe that grokking effect can have better meanings in this large scale, highly nonlinear experiments. And the concept of solution sparsity should be discussed in this large domains, on how it should be “extended”.", "other_strengths_and_weaknesses": "Please refer to other sections of this review. I respect the author’s effort and work for providing theoretical justification in this work. However, I am not still convinced on the logical flow of this manuscript.", "other_comments_or_suggestions": "- In line 432, LSMT→LSTM.", "questions_for_authors": "- Since grokking phenomenon is firstly observed in Transformer network and imposes significance the most in these types of network, why the authors only have chosen to demonstrate for MLP and LSTM networks?", "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": "wpdrQMR5PK", "reviewer": "Reviewer_Cifa", "summary": "This paper studies grokking phenomenons in the setting where the model has a certain special property $P$, and reveals that by using GD with a small but non-zero regularization of $P$ it is possible to observe grokking. In addition, it shows that modifying model depth or performing data selection can also amplify grokking.\n\n### Update after rebuttal\nMy original concern still remains while the authors did not provide further reply. Hence I keep my score unchanged.", "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 thoughtful feedback and the time spent reviewing our work. Our responses to the comments are provided below.\n\n**Weaknesses**\n1) We have rewritten both sections to emphasize the fundamental difference between the two main theorems (2.1 about sparse recovery and 3.1 about matrix factorization) with proof sketches. We first show (using standard results from compressed sensing literature) that, given a matrix $\\tilde{\\mathbf{X}} \\in \\mathbb{R}^{N \\times n}$ and a vector $\\mathbf{b}^* \\in \\mathbb{R}^{n}$ of sparsity $s$ ($\\|\\mathbf{b}^*\\|_0 \\le s$), we can write, for any vector $\\mathbf{b} \\in \\mathbb{R}^{n}$, $\\|\\mathbf{b} - \\mathbf{b}^*\\|_2 \\le C_1 \\| \\tilde{\\mathbf{X}} \\left(\\mathbf{b} - \\mathbf{b}^*\\right)\\|_2 + C_2 | \\|\\mathbf{b}\\|_1 - \\|\\mathbf{b}^*\\|_1|$ where $C_1$ and $C_2$ are constants that depend only on $s$ and $\\tilde{\\mathbf{X}}$. After a short amortization phase,  the term $\\| \\tilde{\\mathbf{X}} \\left(\\mathbf{b}^{(t)} - \\mathbf{b}^*\\right)\\|_2$ vanish (or becomes proportional to the noise if there is any), since we have $ \\tilde{\\mathbf{X}} \\mathbf{b}^{(t)} \\approx \\mathbf{y}^* = \\tilde{\\mathbf{X}} \\mathbf{b}^* + \\xi$. Then comes a second phase, during which $\\|\\mathbf{b}^{(t)}\\|_1$ converges to $\\|\\mathbf{b}^*\\|_1$: our contribution is to show that this phase takes a time inversely proportional to $\\alpha$ (the learning step) and $\\beta_1$ (the regularization strength). \n    \n  The same type of decomposition applies to matrix factorization, with a slight difficulty because not only singular values but also singular vectors are taken into account.  In fact, although the two results (sparse recovery and matrix factorization) are similar, the distinction lies at several levels:\n * in sparse recovery problems, we take into account only the iterate $\\mathbf{b}^{(t)}$ (especially its $\\ell_1$ norm), whereas, in matrix factorization, we take into account not only the singular value matrix $\\Sigma^{(t)}$ (especially its $\\ell_1$ norm) of the iterate $\\mathbf{A}^{(t)} =  \\mathbf{U}^{(t)} \\Sigma^{(t)} \\mathbf{V}^{(t)\\top}$, but also its singular vectors $\\mathbf{U}^{(t)}$ (left) and $\\mathbf{V}^{(t)}$ (right), which makes proofs in the second case slightly more difficult.\n  * if we take a matrix factorization problem and just optimize it with $\\ell_1$, there's no grokking unless the matrix is extremely sparse so that the notion of sparsity prevails over the notion of rank, which shows the point of studying $\\ell_*$ (nuclear norm regularization) separately.\n   * models (linear or not) trained with $\\ell_2$, $\\ell_1$ and $\\ell_*$ have very different properties: the aim of our work is to show that in all these cases, if we want to have a model that generalizes, and which has the desired property (weights of low $\\ell_2$ norm, sparse, or of low rank), then it's better to use a property-specific, weak regularization, and to train the model long enough.\n\nThis extends Luy et al.'s work, which focuses solely on $\\ell_2$. Indeed, there will be no grokking by using only $\\ell_2$ in the problems we study theoretically in the paper. This also shows that the mechanism behind grokking goes beyond the simple $\\ell_2$ norm (we show that with other regularization, we can change the delay between memorization and generalization).\n\n2) Our results cannot be obtained directly from Luy et al.'s work. They show that large-scale initialization and non-zero weight decay lead to grokking provided the model is homogenous. However, as we show, using only $\\ell_2$ (with large-scale initialization or not) in the problems we study theoretically in the paper, there will be no grokking. With large-scale initialization and $\\ell_2$ only, there is an abrupt transition in the generalization error during training, driven by changes in the $\\ell_2$-norm of the model parameters. This transition, however, does not result in convergence to an optimal solution. We call this phenomenon \"grokking without understanding,\" and we attribute it to the fact that the assumptions underlying the theoretical predictions of Luy et al. are violated in our setting.\n\n**Questions: Difference between the core idea of the paper and that of Lyu et al. (2023)** \n\nLuy et al.'s focus on $\\ell_2$ and large initialization. But in our case, there will be no grokking by using only $\\ell_2$ in the problems we study theoretically. So, we extend their work to other regularizers, and there is no need for large-scale initialization in our setting. In addition, we show how grokking can be necessary in specific contexts. Suppose we want a model with low-rank weights and opt to achieve that through $\\ell_*$ regularization. In that case, our work suggests using the lowest possible regularization coefficient (the one the computational resources allow) and training the model as long as possible, far beyond the overfitting point (grokking). We also show that grokking can be extremely amplified/reduced by selecting the data appropriately.", "claims_and_evidence": "Claims are supported by both theoretical and empirical evidence.", "methods_and_evaluation_criteria": "Not applicable, as this paper does not propose new methods.", "theoretical_claims": "I checked the proofs for Section 2 roughly, which appear to be correct.", "experimental_designs_or_analyses": "The experimental designs are sound.", "relation_to_broader_scientific_literature": "The contributions are an addition to the theoretical understanding of the grokking phenomenon.", "essential_references_not_discussed": "The discussion of the essential references are sufficient.", "other_strengths_and_weaknesses": "### Strengths\n\n1. The notations and introduction to the related work (e.g., those about the backgrounds of the problem setting) are satisfying.\n2. Overall the writing is clear. And the proofs are clear and complete.\n3. The findings for the grokking phenomenons in these new settings are also interesting.\n\n\n### Weaknesses\n\n1. This paper gives me a very mixed feeling: I can see that the authors have made significant efforts to try to derive a general result to understand the grokking phenomenon for a certain kind of problem (GD with small regularization), while the authors decided to have two separate but very repetitive sections: the results for Theorem 2.1 and Theorem 3.1 are in fact very similar even though they are studying two different problems. \n\n    This strongly suggests that an underlying mechanism is behind somewhere: it seems that the grokking phenomenons in these two settings are both caused by the different generalization properties (implicit bias) of the early solution and that of the late solution. This is basically the idea proposed by Lyu et al., 2023, although Lyu et al., 2023 focused on classification problem. \n\n2. In my view, the formulations of the regularization are nothing but the types of tendency towards specific late solution, especially when you are using small regularization coefficient, thus can be covered by the core idea of Lyu et al., 2023. In this way, I think the authors should sharpen their message from a higher point of view, clearly summarize their insights compared to those in Lyu et al., 2023, discuss their additional insights at the start of the main result, and reduce the repetitive discussion for the sparse recovery problem and matrix factorization problem. \n\nReference \n\nLyu et al. Dichotomy of early and late phase implicit biases can provably induce grokking.", "other_comments_or_suggestions": "Please see Weaknesses.", "questions_for_authors": "What are the difference between the core idea in this paper and that of Lyu et al., 2023? By saying core idea, I mean the grokking phenomenon is formed by the transition from the early solution with poor generalization to the late solution solution that generalizes better.", "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": "PA4VTyxtCX", "reviewer": "Reviewer_zx76", "summary": "The paper studies grokking i.e., delayed generalization analytically in sparse recovery and matrix factorization settings. The main finding i in the paper provides settings that question the  accepted wisdom in grokking - a small low $L_{2}$ regularization (weight decay) is necessary for grokking and generalization. In other words, $L_{2}$ norm cannot be used as an indicator of delayed generalization. The paper provides extensive analysis and some empirical results to support the main claim\n\n## update after rebuttal\n\nMy current rating is 3. I like the paper but am unwilling to move the score up after reading the rebuttal and other reviews. I have started a discussion with the AC and other reviewers to obtain clarifications but yet to hear from them. Score is unchanged but I have already marked it as a 3.", "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 are grateful to the reviewer for their time and thoughtful feedback. We sincerely appreciate the effort put into evaluating our work and offering constructive suggestions. Below, we respond to each of the points raised.\n\n**Theoretical Claims**\n\nWe've rewritten the proofs of the two main theorems (2.1 and 3.1) more simply, this time with a proof sketch, which we can summarize as follows. The idea is to first show that, given a matrix $\\tilde{\\mathbf{X}} \\in \\mathbb{R}^{N \\times n}$ and a vector $\\mathbf{b}^* \\in \\mathbb{R}^{n}$ of sparsity $s$ ($\\|\\mathbf{b}^*\\|_0 \\le s$), we can write, for any vector $\\mathbf{b} \\in \\mathbb{R}^{n}$, $\n\\|\\mathbf{b} - \\mathbf{b}^*\\|_2 \\le C_1 \\| \\tilde{\\mathbf{X}} \\left(\\mathbf{b} - \\mathbf{b}^*\\right)\\|_2 + C_2 | \\|\\mathbf{b}\\|_1 - \\|\\mathbf{b}^*\\|_1|$ where $C_1$ and $C_2$ are constants that depend only on $s$ and $\\tilde{\\mathbf{X}}$ (notably on the restricted isometry constant of $\\tilde{\\mathbf{X}}$). With this decomposition in mind, we then show that :\n* there is a memorization phase, after which the term $\\| \\tilde{\\mathbf{X}} \\left(\\mathbf{b}^{(t)} - \\mathbf{b}^*\\right)\\|_2$ vanish (or becomes proportional to the noise if there is any), since we have $ \\tilde{\\mathbf{X}} \\mathbf{b}^{(t)} \\approx \\mathbf{y}^* = \\tilde{\\mathbf{X}} \\mathbf{b}^* + \\xi$.\n* and a second phase, during which $\\|\\mathbf{b}^{(t)}\\|_1$ converges to $\\|\\mathbf{b}^*\\|_1$: we show that this phase takes a time inversely proportional to $\\alpha$ (the learning step) and $\\beta_1$ (the regularization strength).\n\nWe arrive at the final result by combining these two steps. The same type of decomposition applies to matrix factorization, with a slight difficulty because not only singular values but also singular vectors are taken into account.\n\n**Weaknesses: Rewriting**\n\nThe rewriting of the proofs allows us to considerably reduce the number of pages in the appendix and to move many results from the appendix to the main text. \n\n**Validation plot** \n\nIn our case, for an iterate $\\mathbf{b}$, the training error is $\\|\\tilde{\\mathbf{X}} \\mathbf{b} - \\mathbf{y}^*\\|_2^2$, and the test (or validation) error is $\\|\\mathbf{b} - \\mathbf{b}^*\\|_2^2$. Normally, the generalization error should be $\\mathcal{E}(\\mathbf{b}) = \\mathbb{E}\\left(\\mathbf{x}^\\top \\mathbf{b} - \\mathbf{y}^*(\\mathbf{x})\\right)^2>$ (expectation with respect to $\\mathbf{x}$ and $\\xi$), where $\\mathbf{y}^*(\\mathbf{x}) = \\mathbf{x}^\\top \\mathbf{b}^* + \\xi$. Assuming $\\mathbb{E}\\xi = 0$, and using $\\Sigma = \\mathbb{E}[\\mathbf{x} \\mathbf{x}^\\top]$, we get $\\mathcal{E}(\\mathbf{b}) = \\left(\\mathbf{b} - \\mathbf{b}^* \\right)^\\top \\Sigma \\left(\\mathbf{b} - \\mathbf{b}^* \\right) + \\mathbb{E}\\xi^2$.\n\nAs we illustrate in the appendix, random matrices, notably Gaussian and Bernoulli with independent entries, allow for the lowest restricted isometry constant and, thus, better recovery of sparse vectors. Under this iid assumption, we have $\\Sigma = \\sigma^2\\mathbf{I}_n$ for a certain $\\sigma > 0$, which implies $\\mathcal{E}(\\mathbf{b}) = \\sigma^2 \\|\\mathbf{b} - \\mathbf{b}^*\\|_2^2 + \\mathbb{E}\\xi^2$. So $\\|\\mathbf{b} - \\mathbf{b}^*\\|_2^2$ captures well the notion of test (or validation) error commonly used in the context of grokking, while $\\mathbb{E}\\xi^2$ capture the irreducible part of that error (we exclude it in the figure for simplicity).\n\n**Could the authors simplify Figure 1 and perhaps Figure 3 to relate these terms to the terms plotted in the figure?**\n\nThanks for the suggestion, we will simplify the figures as best we can.", "claims_and_evidence": "The main claim is that if a problem as a certain property $P$, then appropriate regularization $L_{p}$ is required to reach the optimal solution.\n- Analytical results appear to support the claim\n- Experiments on synthetic data supports this claim as well in Section 2 and Section 3\n- Non-linear teacher student model also appears to support the claim\n\nThe support provided with PINNs and classification are not as clear as the above primarily because the nature of the optimal solution is not obvious (at least to this reader)", "methods_and_evaluation_criteria": "The paper is theoretical work, so the datasets used with nonlinear models are acceptable", "theoretical_claims": "The paper provides a theorem in Section 2 and one in Section 3 for sparse recovery and matrix factorization respectively. The theorems actually summarize several theorems proposed and proved in the appendix. The issue is its easy for the reader to get lost between the claims in Theorem 2.1 (and 3.1) as the text indicates they summarize several claims and proofs in the appendix. It would be better if the Theorem in the main text has a proof sketch in the main paper so that the interested reader can fill in the details from proof(s) in the appendix.", "experimental_designs_or_analyses": "The paper is theoretical work, so the experimental implementation is acceptable", "relation_to_broader_scientific_literature": "The paper does an excellent job relating their work to broader scientific literature", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "# Strengths\n\n- The paper uses a setup that's simple enough for analysis and is able to show that accepted wisdom in grokking, i.e., L_2 norm of parameters may not suffice. Note that this reminds me of a string of papers that questions \"deep learning may not be explained by norms\". Happy to see these papers cited in the work as well\n\n# Weaknesses\n\n- The paper is very dense and some of the sections, if rewritten, would be accessible to a broader audience", "other_comments_or_suggestions": "Grokking is commonly thought of as delayed generalization. Typically papers include a \"training\" vs \"validation\" plot to show delayed generalization. \n\nCould the authors simplify Figure 1 and perhaps Figure 3 to relate these terms to the terms plotted in the figure?\nCould the authors simplify Figure 2 and Figure 4 as there are too many terms plotted in a single plot?", "questions_for_authors": "Please see my comments on including proof sketches for main theorems as well as improving figures in the paper made above.", "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 convincingly demonstrate—both analytically in sparse recovery / matrix‑factorization settings and empirically in deeper non‑linear networks—that delayed generalization (grokking) is not tied to the Euclidean (L2) norm but rather to whichever structural property R is truly capacity‑controlling for the task (sparsity, low rank, etc.).  Reviewers agreed the theoretical results are sound, the revised draft now contains clear proof sketches and cleaner figures, and extensive synthetic‐and‑real experiments—including teacher–student MLPs, PINNs, and LSTMs—validate the predicted scaling and show that depth‑induced implicit bias or targeted data selection can amplify or suppress the effect.

Although the paper remains dense, the authors’ rewrite addresses presentation concerns: each main theorem is now better motivated in the body, repetitive material has been pruned, and the connection between the linear theory and non‑linear evidence is explicitly framed as an empirical test of the theory’s insight rather than a claimed formal extension. 

In addition, even though the reviewers were concerned that the results are limited to linear models, I believe this is acceptable for two reasons. The goal of the paper is not to make claims about grokking in all model families, but rather to isolate and analyze the underlying mechanism in a setting that is amenable to precise analysis. Second, the empirical evidence presented across a range of non-linear architectures (MLPs, PINNs, LSTMs) suggests that the insights obtained in the linear case continue to hold qualitatively.
Decision: Accept (poster)