Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking
Authors: Kaifeng Lyu
Jikai Jin
Zhiyuan Li
Simon Shaolei Du
Jason D. Lee
Wei Hu
Venue: ICLR 2024 poster (2024)
Keywords: grokking
implicit bias
margin
kernel
training dynamics
generalization
Primary area: learning theory
TL;DR: Due to a dichotomy of early and late phase implicit biases, training homogeneous neural nets with large initialization and small weight decay provably induces grokking.

Abstract: Recent work by Power et al. (2022) highlighted a surprising "grokking" phenomenon in learning arithmetic tasks: a neural net first "memorizes" the training set, resulting in perfect training accuracy but near-random test accuracy, and after training for sufficiently longer, it suddenly transitions to perfect test accuracy. This paper studies the grokking phenomenon in theoretical setups and shows that it can be induced by a dichotomy of early and late phase implicit biases. Specifically, when training homogeneous neural nets with large initialization and small weight decay on both classification and regression tasks, we prove that the training process gets trapped at a solution corresponding to a kernel predictor for a long time, and then a very sharp transition to min-norm/max-margin predictors occurs, leading to a dramatic change in test accuracy.

Published as a conference paper at ICLR 2024
DICHOTOMY OF EARLY AND LATE PHASE IMPLICIT BIASES CAN PROVABLY INDUCE GROKKING
Kaifeng Lyu"
Jikai Jin'
Zhiyuan Li
Princeton University
Stanford University
Toyota Technological Institute at Chicago dilyyanmi8tticn
kiyu@cs princeton..... edu
Simon S. Du University of Washington ssdu@cs. washington.een edu
Jason D. Lee
Wei Hu
Princcton University
University of Michigan
jasonlee@princetton.
vvh aen@@miice umich. edu
ABSTRACT
Recent work by Power al. (2022) highlighted a surprising "grokking" phe- nomenon in learning arithmetic tasks: a neural net first "memorizes" the training set, resulting in perfect training accuracy but near-random test accuracy, and af- ter training for sufficiently longer, it suddenly transitions to perfect test accuracy. This paper studies the grokking phenomenon in theoretical setups and shows that it can be induced by a dichotomy of early and late phase implicit biases. Specif- ically, when training homogeneous neural nets with large initialization and small weight decay on both classification and regression tasks, we prove that the training process gets trapped at a solution corresponding to kernel predictor for a long time, and then very sharp transition to aaan--oommmmx-maaeee predictors occurs, leading to a dramatic change in test accuracy.
INTRODUCTION
The generalization behavior of modern over-parameterized neural nets has been puzzling: these nets have the capacity to overfit the training set, and yet they frequently exhibit a small gap between training and test performance when trained by popular gradient-based optimizers. A common view now is that the network architectures and training pipelines can automatically induce regularization effects to avoid or mitigate overfitting throughout the training trajectory.
Recently, Power et (2022) discovered an even more perplexing generalization phenomenon called grokking: when training neural net to learn modular arithmetic operations, it first "memorizes" the training set with zero training error and near-random test error, and then training for much longer leads to sharp transition from generalization perfect generalization. See Section 2 for our production this phenomenon. Beyond modular arithmetic, grokking has been reported in learning group operations (Chughtai al., 2023), learning sparse parity (Barak al., 2022; Bhattamishra et al., 2023), learning greatest common divisor (Charton, 2024), and image classification (Liu et al., 2023; Radhakrishnan et al., 2022).
Different viewpoints on the mechanism of grokking have been proposed, including the slingshot mechanism (cyclic phase transitions) (Thilak al., 2022), random walk among minimizers (Mil- lidge, 2022), slow formulation of good representations (Liu et al., 2022), the scale of initializa- tion (Liu et al., 2023), and the simplicity of the generalizable solution (Nanda al., 2023; Varma et al., 2023). However, existing studies failed to address two crucial aspects for gaining a compre- hensive understanding grokking:
1. No prior work has rigorously proved grokking neural network setting.
2. No prior work has provided quantitative explanation to why the transition from memoriza- tion to generalization is often sharp, instead of gradual.
"Equal Contribution
Code available htc.ge///gthh....cc///aaaaaaaaa/ chottmmm

---

Published as a conference paper at ICLR 2024
train
test
train
test
100%
100%
90%
90%
accuracy
accuracy
80%
80%
70%
70%
60%
60%
50%
50%
10
10:
10a
10a
10*
109
10
100
10a
10a
10a
10*
10a
20'
step
125
125
0.6
L parameter morm
100
100
10s5
75
1ocs
0.6
75
pararme parameter
50
50
0.2
25
25
0.0
0.0
0
100
10
10?
10
10a
10
10
010
10a
10²
102
10a
10a
10t
107
step
step
(a) Grokking
(b) "Misgrokking"
Figure 1: Training two-layer diagonal linear nets for linear classification can exhibit (a) grokking, for a dataset that can be linearly separated by a 3-sparse weight vector, or (b) "misgrokking", for dataset that can linearly separated by large T,2-margin. See Section 3.2.3 for details.
Our Contributions. In this work, we address these limitations by identifying simple yet insightful theoretical setups where grokking with sharp transition can be rigorously proved and its mechanism can be intuitively understood. Our main intuition is that optimizers may implicitly induce different biases in early and late phases; grokking happens if the early phase bias implies an overfitting solution and late phase bias implies generalizable solution.
More specifically, we focus on neural nets with large initialization and small weight decay. This is inspired by a recent work (Liu et al., 2023) showing that training with these two tricks can in- duce grokking on many tasks, even beyond learning modular arithmetic. Our theoretical analysis attributes this to the following dichotomy of early and late phase implicit biases: the large initializa- tion induces very strong early phase implicit bias towards kernel predictors, but over time, it decays and competes with late phase implicit bias towards min-norm/mon -argin predictors induced by the small weight decay, resulting in a transition that turns out to be provably sharp in between.
This implicit bias result holds for homogeneous neural nets, a broad class of neural nets that in- clude commonly used MLPs and CNNs with homogeneous activation. We further exemplify this dichotomy in sparse linear classification with diagonal linear nets and low-rank matrix completion with over-parameterized models, where we provide theory and experiments showing that the ker- nel predictor does not generalize well but the min-norm/maxx-arye maagmnn predictors do, hence grokking happens (see Figures la and 3).
Based on these theoretical insights, we are able construct examples where the early phase implicit bias leads to a generalizable solution, and the late phase bias leads to overfitting. This gives rise to new phenomenon which we call "misgrokking"": neural net first achieves perfect training and test accuracies, but training for longer time leads to a sudden big drop in test accuracy. See Figure 1b.
Our proof of the kernel regime in the early phase is related the Neural Tangent Kernel (NTK) regime studied in the literature (Jacot et al., 2018; Chizat et al., 2019). However, our result is qualitatively different from existing NTK analyses significantly (due to weight decay), while in the usual NTK regime, the weight changes only by small amount, and so does its norm. Our proof relies on careful analysis of the norm and direction of the weight, which enables us obtain a tight bound on the time spent in the kernel regime. This novel result of kernel regime in homogeneous neural nets under weight decay may be of independent interest. We further analyze the late phase using gradient convergence argument, and it turns out that slightly longer time suffices for convergence to the max margin/min norm predictor, thus exhibiting sharp transition between the phases.
that in our setting, the weight norm changes
2

---

Published as a conference paper at ICLR 2024
A-o A.0.000
0.0000
0.003
Irain
Fal
Seat
1.0%
300%
100%
10"
30%
93%
80%
50%
20%
23%
20%
101
(a) He init, 10-4
(b) Effect of Initialization Scale
(c) Effect of Weight Decay
Figure 2: Training two-layer ReLU nets for modular addition exhibits grokking with He initializa- tion and weight decay. Enlarging the initialization scale or reducing the weight decay delays the sharp transition in test accuracy.
Concretely, our contributions can be summarized as follows:
1. For training homogeneous neural nets with large initialization and small weight decay on clas- sification and regression tasks, we prove a very sharp transition from optimal solutions of kernel SVM/rerression to KKT solutions of global max maagim/rgn//mmnnooo problems associated with the neural net.
2. For classification and regression tasks respectively, we provide concrete examples with diagonal linear nets and overparameterized matrix completion, showing either grokking or misgrokking.
2 MOTIVATING EXPERIMENT. GROKKING IN MODULAR ADDITION
In this section, we provide experiments to show that the initialization scale and weight decay are important factors in the grokking phenomenon, in the sense that enlarging the initialization scale or reducing the weight decay delays the sharp transition in test accuracy.
Grokking in Modular Addition. Following many previous works (Nanda et al., 2023; Gromov, 2023), we focus on modular addition, which belongs to the modular arithmetic tasks where the grokking phenomenon was first observed (Power et al., 2022). The task is to learn the addition operation over Ep: randomly split {(a,b, b, a + b = e (mod p}} into training and test sets, and train a neural net on the training set to predict given input pair (a. The grokking phenomenon has been observed in learning this task under many training regimes, including training transformers and MLPs on cross-entropy and squared loss with various optimizers (full-batch GD, SGD, Adam, Adam W, etc.) (Power et al., 2022; Liu et al., 2022; Thilak et al., 2022; Gromov, 2023). For sim- plicity, focus on simple setting without involving too many confounding factors: training two-layer ReLU net with full-batch GD. More specifically, we represent the input (a, b) E Ep x Zp as the concatenation of the one-hot representations of a and b, resulting in an input vector x E R and then the neural net processes the input as follows:
((0: = W2 ReI.U(W b1),
where 0 = (Wi, bi, W2) is the parameter vector consisting of Wi BR E RA W2 Rxxh h is the width, and ReLU(v) stands for the ReLU activation that maps each entry is of a vector max(u;, 0}. The neural net output f(0; x) 12" is treated logits and fed into the standard cross-entropy loss in training. We set the modulus p as 97, the training data size as 40% of the total number of data (p²), the width h as 1024, the learning rate as 0.002 and weight decay as 10-4 Figure 2a successfully reproduces the grokking phenomenon: the training accuracy reaches 100% in 10 steps while the test accuracy remains close to random guessing, but after training for one additional order of magnitude in the number of steps, the test accuracy suddenly jumps to 100%.
natural question is how this grokking phenomenon depends on the training tricks in the pipeline. Liu et al. (2023) showed that, empirically, using large initialization scale and small but non- zero weight decay can induce grokking on various tasks even beyond modular arithmetic, including image classification MNIST and sentiment classification IMDB. Indeed, experiments on modular addition can confirm the importance of initialization scale and weight decay.
Effect of Initialization Scale. To study the effect of initialization, we scale up the standard He initialization (He et al., 2015) by a factor of a > 0, and then run the same training procedure. However, the model at large initialization produces very large random number for each training
3

---

Published as a conference paper at ICLR 2024
sample, which induces training instability. To remove this confounding factor, we set the weights W2 in the second layer to 0 so that the initial model output is always O. Results in Figure 2b show that the sharp transition in test accuracy is delayed when a increases from 1 to 100.
Effeet of Weight Decay. In Figure 2c, we run experiments with different weight decay values. We observe that increasing weight decay makes the transition happen earlier. However, when weight decay surpasses a small threshold, the regularization strength can become so strong that it even hinders the optimization of the training loss and causes the training accuracy to collapse eventually. This collapse due to weight decay is also consistent with the observation in Lewkowycz. and Gur-Ari (2020). Conversely, as weight decay decreases, the grokking phenomenon becomes more significant because now it takes more time for the test accuracy to exhibit the sharp transition. But the weight decay has to be non-zero: we will discuss in Appendix B. that training without weight decay leads to perfect generalization in the end, but the transition in test accuracy is no longer sharp.
Explanation? Liu et al. (2023) built their intuition upon an empirical observation from Fort and Scherlis (2019): there is narrow range of weight norm, called Goldilocks zone, where the gen- eralization inside the zone is better than the outside. It is argued that the weight norm can take a longer time to enter this Goldilocks zone when training with a larger initial weight norm or smaller weight decay. hence causing the delayed generalization. However, Fort and Scherlis (2019); Liu et al. (2023) did not provide much explanation on why such narrow Goldilocks zone could exist in the first place. When talking about MIPs and CNNs with ReLU activation, which are commonly used in practice, this becomes even more mysterious: the output functions f(0; æ:) of these neural nets are homogeneous to their parameters, i.c. f(c0; a) = c² f(0; x) for all > 0 (Lyu and Li, 2020; Ji and Telgarsky, 2020a; Kunin et al., 2023). As a result, is impossible to explain the effect of norm just from the expressive power, since all classifiers that can be represented with a certain norm can also be represented by all other norms. This motivates us to dive deep into the training dynamics in grokking via rigorous theoretical analysis.
GROKKING WITH LARGE INITIALIZATION AND SMALL WEIGHT DECAY
In this section, we present our theory on homogeneous neural nets with large initialization and small weight decay, which attributes grokking to a dichotomy of early and late phase implicit biases.
3.1 THEORETICAL SETUP
We focus on training models for classification and regression. Let Dx be the input distribution. For every input
Dx, let y (xx) be the classification/rerrsssinn target of æ. We parameterize the model with 0 and use f(0; x) to denote the model output on input æ. We train the model by optimizing an empirical loss on dataset {ax; y:) where the data points are drawn i.i.d. as
2; ~ >>xx yi = y² (x::). For each i E [m)] we write (;(0) == [(0; æx) for short. We assume the model is homogeneous with respect to its parameters 0, common assumption in analyzing neural networks (Lyu and Li, 2020; Ji and Telgarsky. 2020a: Telgarsky, 2023; Woodworth et al., 2020).
Assumption 3.1 For all æ. f (0; is L-homogeneous (ie., [c( ct fc(0) for all 0) and C² smooth with respect to O.
The two-layer net in our motivating experiment and other MLPs and CNNs with ReLU activation are indeed L-homogeneous. However, we also need to assume the C²-smoothness to ease the definition and analysis of gradient flow, which excludes the use of any non-smooth activation. This can be potentially addressed by defining and analyzing gradient flow via Clarke's 's subdifferential (Clarke, 1975; 2001; Lyu and Li, 2020; and Telgarsky, 2020a), but we choose to make this assumption for simplicity. Note that our concrete examples of grokking on sparse linear classification and matrix completion indeed satisfy this assumption.
As motivated in Section 2, we consider training with large initialization and small weight decay. Mathematically, we start training from adi where Onni E Rd is a fixed vector and a is a large factor controlling the initialization scale. We also assume small but non-zero weight decay À > 0. To study the asymptotics more conveniently, we regard À as a function of a with order X(a) = ((a P) for some positive p= O(1)
To avoid loss explosion at initialization, we deterministically set or randomly sample Oinit in a way that the initial output of the model is zero. This can be done by setting the last layer weights to zero
4

---

Published as a conference paper at ICLR 2024
(same experiments in Section 2), using the symmetrized initialization in Chizat et al. (2019), or using the "difference trick" in Hu et al. (2020). This zero-output initialization is also used in many previous studies of implicit bias (Chizat et al., 2019; Woodworth et al., 2020; Moroshko al., 2020) for the sake of mathematical simplicity. We note that our analysis should be easily extended to random initialization with small outputs, just in Chizat al. (2019); Arora et al. (2019b).
Assumption 3.2 f f(0iii: x) 0 for all æ.
We define the vanilla training loss the average over the loss values of each individual sample:
(())
((f(()::.)
(2)
With weight decay X, it becomes the following regularized training loss:
X 1012
(x ( +
(3)
For simplicity, in this paper, we consider minimizing the loss via gradient flow de = -Lxx
dr which is the continuous counterpart of gradient descent and stochastic gradient descent when the learning rate goes to O. We use (tt; 0o) to denote the parameter at time t by running gradient flow from og and we write (t) when the initialization On is clear from the context.
3.2 GROKKING IN CLASSIFICATION
In this subsection, study the grokking phenomenon on classification tasks. For simplicity, restrict to binary classification problems, where y (a:) E [-1,+1]. While in practice it is standard to use the logistic loss {Øū y) log(1 + + (a.k.a. binary cross-entropy loss), we follow the implicit bias literature to analyze the exponential loss 8) = e 29 as simple surrogate, which has the same tail behavior as the logistic loss when yú + +oc and thus usually leads the same implicit bias (Soudry et al., 2018; Nacson et al., 9b; Lyu and Li, 2020; Chizat and Bach, 2020),
By carefully analyzing the training dynamics, we rigorously prove that there is sharp transition around x log Before this point, gradient flow fits the training data perfectly while maintaining the parameter direction within a local region near the initial direction, which causes the classifier to behave like kernel classifier based on Neural Tangent Kernel (NTK) (Jacot et al., 2018; Arora et al., 2019b;c). After the transition, however, the gradient flow escapes the local region and makes efforts to maximize the margin. Following the nomenclature in (Moroshko al., 2020), we call the first regime the kernel regime and the second regime the rich regime. The key difference to the existing works analyzing the kernel and rich regimes (Moroshko al., 2020; Telgarsky, 2023) is that the transition in our case is provably sharp, which a crucial ingredient for the grokking phenomenon.
3.2.1 KERNEL REGIME
First, we show that gradient flow gets stuck in the kernel regime over the initial period of ميs log a, where the model behaves as if it were optimizing all x, we define f x) as the NTK feature of æ. As we are considering over-parameterizze models, where the dimension of O is larger than the number of data n, is natural to assume that the NTK features of the training data are linearly separable (Ji and Telgarsky, 2020b; Telgarsky, 2023):
linearized model. More specifically, for
Assumption 3.3 There exists h such that yi (V (Omitt) h2 f.ror all i E [n].
Let hitk be the unit vector so that a linear classifier with weight hitk can attain the max L -m-minn on the NTK features of the training data, (O He )z and Yatk be the corresponding -²_ margin. That is, his is the unique unit vector that points to the direction of the unique optimal solution to the following constrained optimization problem:
min
1111112 s.t. y: (V h) 1, Vie [n],
(KI)
and 7atk == minke()) (u. (V (00m)) h h.n)))
The following theorem states that for any e E (0, 1), the solution found by gradient flow at time 1-c log a represents the same classifier as the max I - margin linear classifier on the NTK features:
5

---

Published as a conference paper at ICLR 2024
Theorem 3.4 For any all constants (0.1). letting Te (a) > log it holds that
Va IRd
a- ((a) [(O(T, (a); al mii)) (V x)) kmik)
lim
where 7(0)=
log a X is normalizing factor.
Yatk
Key Proof Insight. The standard NTK-based dynamical analysis requires (tt; amitt) z af but the weight decay in our case can significantly change the norm, hence violating this condi- tion. The key ingredient of our proof is a much more careful dynamical analysis showing that (t: aQitt) = ae 0 ie., the norm of ((: min)) decays but the direction remains close to Ount. Then by L-homogeneity and Taylor expansion, we have
f(0: (are - f( 0:))) z (ax -r)= (f(Omiti + (V mits 2:),
Oinit
= - LAM (( x)),
0mit)
The large scaling factor -Ltt above enables fitting the dataset even when the change in direction θ mit is very small. Indeed, we show that gxe Omit ~ =-_aa Lake LM log a" hink at time Te log a by closely tracking the dynamics, hence completing the proof. See Appendix C.2 for details. When a e LAt is no longer a large scaling factor, namely t = (log a + w(1),, this analysis breaks and the dynamics enter the rich regime.
Patk
3.2.2 RICH REGIME
Next, we show that continuing the gradient flow for a slightly long time to time 1+c log it is able to escape the kernel regime. More specifically, consider the following constrained optimization prob- lem that aims at maximizing the margin of the predictor f(0; x) the training data {( M
min
21 1013
s.t. (0) 1, vi E [n].
(RI)
Then have the following directional convergence result: Theorem 3.5 For all constants > O and for every sequence fas h221 with ag - +oc, letting T (a) == + log a, there exists a time sequence {t*}az) such that } log as I tv S T (ax) and every limit point of alta (rnin)
2 =1} is along the direction of a KKT point of (R1).
For the problem (R1), the KKT condition is a first-order necessary condition for global optimality, but it is not sufficient in general since fi) can be highly non-convex. Nonetheless, since gradient flow can easily get trapped spurious minima with only first-order information, the KKT condition is widely adopted as surrogate for global optimality in theoretical analysis (Lyu and Li, 2020; Wang et al., 2021; Kunin al., 2023).
Key Proof Insight. The key is to use the loss convergence and norm decay bounds from the kernel regime as a starting point to obtain small upper bound for the gradient (O). Then we connect the gradient upper bounds with KKT conditions. See Appendix C.3 for the proof.
3.2.3 EXAMPLE: LINEAR CLASSIFICATION WITH DIAGONAL LINEAR NETS
Now, we exemplify how our implicit bias results imply a sharp transition in test accuracy in a concrete setting: training diagonal linear for linear classification. Let Dx be distribution over Rd and y" (x) = aig(((",,0 æ)) {+1} be the ground-truth target for binary classification, where w² E 12d is an unknown vector. Following Moroshko al. (2020), we consider training a so-called two-layer "diagonal" linear net to learn y². On input æ E IRd the model outputs the
following function f00; x)), where : (u, v) Rdd Rdd R is the model parameter.
f(0: x))
Σ
2 Σ Dk
(4)
This model is considered diagonal net because it can be seen pparsely-coneectee two-layer net with 2d hidden neurons, each which only inputs and outputs a single scalar. More specifically, the input is first expanded 2d-dimensional vector (æ, -r). Then the k-th hidden neuron

---

converges to Published as a conference paper at ICLR 2024
in the first half inputs 2& and outputs u,JK,, and the k-thh hidden neuron in the second half inputs - -A and outputs UyTk-- The output neuron in the second layer shares the same weights as the hidden neurons and takes a weighted sum over the hidden neuron outputs according to their weights: k=1
f(0; æ) = Ik=1 UK (ux= rx) + Ek=1 UK (ok- (-rx)). Note that the model output always linear with the input. For convenience, we can write f(0; x) = (w(0), æ) with effective weight (() == u 02 202
In this setup, we can understand the early and late phase implicit biases very concretely. We start training from a large initialization: (0) init. where a > 0 controls the initialization scale and 0 Oniit (1, 1). That is, Mg Uk a for all 1 I k S d at t = 0. As noted in Moroshko et al. (2020), the kernel feature for each data point is (' x) = (2x, -2x), so the max L -mmain kernel classifier is the max -mrminn linear classifier. This leads to the following corollary:
> log a. For all constants c E (0, 1), as a + too, the normalized effective weight vector wn wO (a):a0m))))
Corollary 3.6 (Diagonal linear nets, kernel regime) Let To (a)
[0(T (0))amm)))l
the max --marrnin direction, namely w == arg max (minge(x) Mi (wW xz,))
In contrast, if we train slightly longer from T-' (a) to T:t (a), we can obtain KKT point of (R1).
Corollary 3.7 (Diagonal linear nets, rich regime) Let T (a) == 1 log a. For all constants E (0,1) 1), and for every sequence {ax)k21 with Ock + +00, there exists a time sequence {1x]221 Wlle:an0mmi
2 } is along the direction of a KKT point of the problem min |/aal13 + ||=1|| s.I. y: n02 m.) 2 1.
such that 2 log Okk S 1/ S T (ax) and every limit point of [10tt Orne)) Ila
As noted in these two works, in the diagonal linear case, optimizing (R1) to the global optimum is equivalent to finding the max -ma-inn linear classifier: min ||w||| s.t. y1 (w, 2 1, Vi E [n], or equivalently, wi := arg = (minge[m) (w, ..)}. Therefore, our theory suggests a sharp transition at time X log a from max I, m.a-margin linear classifier to max T, mmrginnn linear classifier.
While this corollary only shows KKT conditions rather than global optimality, we believe that one may be able to obtain the global optimality using insights from existing analysis of training diagonal linear nets without weight decay (Gunasekar et al., 20;8b; Moroshko al., 2020).
Empirical Validation: Grokking. As it is well known that maximizing the L--mmagiin can bet- ter encourage sparsity than maximizing the L'-margin, our theory predicts that the grokking phc- nomenon can be observed in training diagonal linear nets for k-sparse linear classification. To verify this, we specifically consider the following task: sample n data points uniformly from [=1] for a very large d, and let the ground truth be a linear classifier with a k-sparse weight, where the non-zero coordinates are sampled uniformly from {+1}*. Applying standard generalization bounds based on Rademacher complexity can show that the max L'-mmagiin linear classifier needs O(kd) to gener- alize, while the max L -marinn linear classifier only needs O(k2 log d). See Appendix E.2 for the proof. This suggests that the grokking phenomenon can be observed in this setting. In Figure 1a, we run experiments with n = 256, d = 10 A = 3, large initialization with initial parameter norm 128, and small weight decay À 0.001. As expected, we observe the grokking phenomenon: in the early phase, the net fits the training set very quickly but fails to generalize; after a sharp transition in the late phase, the test accuracy becomes 100%.
Empirical Validation: Misgrokking. From the above example, we can see that grokking can be induced by the mismatch between the early-phase implicit bias and data and good match between the late-phase implicit bias and data. However, this match and mismatch can largely depend on data. Conversely, if we consider the case where the labels in the linear classifica- tion problem are generated by a linear classifier with a large L L'-mangin, the early phase implicit bias can make the neural net generalize easily, but the late phase implicit bias can destroy this good generalization since the ground-truth weight vector may not have large I Im--magin. To justify this, we first sample a unit-norm vector wt, then sample a dataset from the distribution (æ, v) ~ (z + sign((z, w²) ))o",, sign((z, w²)), where
(0 Ia) (i.e., a Gaussian distribu- tion that is separated by a margin in the middle). We take y = 25 and sample n = 32 points. Indeed, observe phenomenon which we call "misgrokking" : the neural net first fits the training set and achieves 100% test accuracy, and then after training for sufficiently longer, the test accuracy drops to nearly 50%. See Figure 1b.
7

---

Published as a conference paper at ICLR 2024
3.3 GROKKING IN REGRESSION
In this subsection, we study the grokking phenomenon on regression tasks with squared loss ((0, y) y)² Paralleling the classification setting, we analyze the behavior of gradient flow in both kernel and rich regimes, and show that there is sharp transition around & log
3.3.1 KERNEL REGIME
Similar to the classification setting, we first show that gradient flow gets stuck in the kernel regime the initial period of X 1- log a time. For all æ we define f 2) as the NTK feature of a. For over- weer-parameteiized models, where the dimension of is larger than the number of data n, we make the following natural assumption that is also widely used in the literature (Du al., 2019b; Chizat et al., 2019; Arora et al., 2019b):
Assumption 3.8 The NTK features of training samples (V (O x)};=1 are linearly independent.
Now we let otk be the vector with minimum norm such that the linear predictor g - (g. h) perfectly fits mization problem:
That is, mt is the solution to the following constrained opti-
min
1 s.t. (VI((mi))))) = Vi E [m)]
(K2)
Then we have the following result that is analogous to Theorem 3.4. See Appendix C.2 for the proof. Theorem 3,9 For all constants c E (0, 1), letting T: (a) == _ log a, it holds that
Va Rdd.
lim f(O(T, (a); mi)) = ((0iit 2)), h a-++0
3.3.2 RICH REGIME
Similar to the classification setting, we then show that gradient flow is able to escape the kernel regime time 1+c log a. Specifically, consider the following constrained optimization problem
X that searches for the parameter with the minimum norm that can perfectly fit the training data:
2 111011 s.t. f((0)
min
Vi [n].
(R2)
Then we have the following convergence result. See Appendix D.2 for the proof. Theorem 3.10 For any constant > letting T (a) == + log a, for any sequence of {axkx21 with Og + +0, there exists a time sequence {fx] satisfying X log as S 4 S (ax)) such that []0(t;: au (r.fmi))l la are uniformly bounded and that every limit point of (t(tt: k 1}
a KKT point of (R2).
3.3.3 EXAMPLE: OVERPARAMETERIAZ MATRIX COMPLETION
As example, consider the matrix completion problem of recovering low-rank symmetric matrix based on partial observations of its entries. Specifically, we let X E ddddd be the ground- truth symmetric matrix with rank r < d and we observe a (random) set entries indexed by Ω C [d] x [d]. In this setting, the input distribution Dx is the uniform distribution on the finite set [d] x [d]. For every input x = (i.j) Dx, y (x) (Pa. X*) where Pa P==ca e.e..
Following previous works (Arora et al.. 2019a; Razin and Cohen, 2020: Woodworth et al., 2020; Li al., 2021), we solve the matrix completion problem using the matrix factorization approach, i.e. we parameterize matrix as W = UUT VVT, where U, V E IBdxd and optimize U, V so that W matches with the ground truth on observed entries. This can be viewed a neural net that parameterized by f = (vec(U), vee( V)) E 2dd and outputs f(0 x) = (P.- W(0)) given æ = (i,j), where W(0) UUT VVT. is easy to that ʃ(0; a:) is 2-homogeneous, satisfying Assumption 3.1. Given a dataset {(x;, yt) ) we consider using gradient flow to minimize the l-rregularized square loss defined in (3) which can be equivalently written as
1 n Σ (f(0; 31)² 2 >
(())
(IUII2 +
(5)
In the kernel regime, Theorem 3.9 implies the following result for identity initialization.
8

---

Published as a conference paper at ICLR 2024
Corollary 3.11 (Matrix completion, kernel regime) Let intt (vec( [). vec(I)). For all con-
stants с E (0,1), letting T- (a) = - log a, it holds for all (i.j) E [d] x [d] that
[ (a); (m)))) - { [Xi +00
(f. E Ω or (1,)
lim
0 otherwise.
In other words, in the kernel regime, while GD is able to fit the observed entries, it always fills the unobserved entries with 0, leading to significant failure in recovering the matrix.
In the rich regime, Theorem 3.10 implies that the solution transitions sharply near 1 log to a KKT point of the norm minimization problem. If this KKT point globally optimal, then it is known that ( is the minimum nuclear norm solution to the matrix completion problem: min || [W || s.t.
W W Wi Xiy, vi..j) E 2 (Ding et al., 2022, Theorem 3.2). is known that the minimum nuclear norm solution can recover the ground truth with very small errors when Od d) entries are observed (Theorem F.5).
Although it is not proved that this KKT point is globally optimal, the following theorem provides positive evidence by showing that gradient flow eventually converges global minima and recovers the ground truth, if the time is not limited around X loga. a.
Theorem 3.12 (Matrix completion, rich regime) Suppose that rank(X') - O(1) and X Vx- Ex- VX. a SVD of X", where each row Vx- has 2e norm bounded by / the mumber of observed entries satisfies N Z 1* dloog d, then any > O, the gradient trajectory (2(t))c2o for (5) starting from random initialization (O)) (V(0));) fnd
N(al = 1) 0²,
converges to global minimizer Oz of Lx. Moreover, [ X' S with probability 2 1-d
d
Remark 3.13 The small random perturbation to the identity initialization (ad, ad) in the above theorem is needed to guarantee that gradient flow does not get stuck at saddle points and thus converges to global minimizers almost surely. By the continuity of the gradient flow trajectory w.E.t initialization, the conclusion of Corollary 3.11 still holds if = a(a) is sufficiently small.
sim- ple algorithmic task: matrix completion for the multipli- cation table. Here, the multiplication table refers to matrix X where the (i.t))- is ij/d?. It is easy to see that this matrix is of rank-1 since X* uu" for
Empirical Validation. We verify our theory
train
test
200
0.10
150
MEEATTTEE nomm
0.06
= (0, 1 2 2: 2: .. d=1). This means the late phase bias can complete the matrix with few observed entries; how- ever, the early phase bias fills every unobserved entry with 0, leading to a large test loss. This dichotomy of implicit biases implies grokking. To check this, we set d = 97 and randomly choose 5% of the entries as the training 10, learn- ing rate = 0.1, weight decay 10-A. Figure 3 shows that grokking indeed happens: GD takes ~ 10a steps
100
0.02
0.00
0
10'
10
10'
10"
1.0"
10'
step
set. We run GD with initialization scale
Figure 3: Matrix completion for the multiplication table exhibits grokking.
minimize the training loss, but the sharp transition in test loss occurs only after ~ 10th steps.
4 CONCLUSION
In this paper, we show that the grokking phenomenon provably occurs in several setups. Our results suggest that the sharp transition in test accuracy may stem from dichotomy of implicit biases between the early and late training phases. While the early phase bias guides the optimizer to overfitting solutions, it quickly corrected when the late phase bias takes effect.
Some limitations of our work are as follows. First, our work only studies the training dynamics with large initialization and small weight decay, but these may not be the only source of the dichotomy of the implicit biases. We further discuss in Appendix B that the late phase implicit bias can also be induced by implicit margin maximization and sharpness reduction, though the transition may not be as sharp the weight decay case. Also, our work focuses understanding the cause of grokking but does not study how to make neural nets generalize without so much delay in time. We leave it to future work explore the other sources of grokking and practical methods to eliminate grokking.
9

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
Kaifeng Lyu partly supported by NSF and ONR. Simon S. Du supported by supported by NSF ISS 21 10170, NSF DMS 2134106, NSF CCF 2212261, NSF ISS 2143493, NSF CCF 2019844, NSF IIS 2229881. Wei Hu is supported by Google Research Scholar Program.
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors. Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 242-252. PMLR, 09-15 Jun 2019. URL https: http://pooeeiiieeee een zhul 9a, html.
8 m.ml r. pre.ss 9t////////////
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix fac- torization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buee E. Fox, and R. Garnett, editors, Advances in Neural Information Prucessing Systems, volume 32. Curran Associates, Inc., 2019a. URL https: httpe//://ppceeeeeeeeeeeee http://poceeeenccceeoece neurips. oc cc/paper....... willes/paper/2000 F1le:/o0075S550007///88111111111
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine- aine-graindd analysis of optimization and generalization for overparameterized two-layer neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 322-332. PMLR, 09-15 Jun 2019b. URL https: http://poceeeennne/ee/e mlr. www.sppessseeeee aroral 9a .html.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In H. Wallach, H. Larochelle, A. Beygelz- imer, F. L'Alché-Buc. E. Fox, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019c. URL UtL.httpss//prrccedinn..oo ips. cc/paper.. 8c4-Paper. pdf.
Pranjal Awasthi, Natalie Frank, and Mehryar Mohri. On the Rademacher complexity of linear hypothesis sets. arXiv preprint arXiv:2007.11045. 2020.
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hid- den progress in deep learning: SGD learns parities near the computational limit. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Informa- fion Processing Systems, volume 35, pages 21750-21764, Curran Associates, Inc., 2022. URL https: httpp//popeeeinnnnncee.eeeee neur neurips... ps wwppppper_aeeeeeeeeeeeeeeeeeeeee files/ patas/aper/22222 ftte:/fleeeee 884bba E65392170763b27c355 7bddee Paper- Conferenceececccccocccoo
Satwik Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability to learn sparse boolean functions. In Proceedings of the 6l Anmual Meeting of the Association Computational Linguistics (Volume 1: Long Papers), pages 5767-5791, 2023.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an Ornstein- Seieenn-Uhlebeckk like process. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 483-513. PMLR, 09-12 Jul 2020. URL eedings. www.rrr..... www...ppeess/11///bbeees html.
Emmanuel J Candes and Yaniv Plan. Matrix completion with noise. Proceedings of the IEEE, 98 (6):925-936. 2010.
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. I'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems. volume 32. Curran Associates, Inc., 2019. URL https: http:///ttttetee.tct/tttooo aper/2019 / f 1 e oec.
wwuripp.....uuu..p. wwwpeeurrssspcpppeeeeeee cc/paper __ieee/pp_pppp
10

---

Published as a conference paper at ICLR 2024
Francois Charton. Learning the greatest common divisor: explaining transformer predictions. In The Tweifth International Conference on Learning Representations, 2024. URL https://o/ enreview net/ wwtpeeeeettttteeeeeeeeeeeeeeeeee forum?1d=cmcDONNP...
Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. Conference on Learning Theory, pages 1305-1338, PMLR, 2020.
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: http://ppooeeeeceeeeeeeeeeceeeet www.ipeuuppp..ee cctpaper_paeaeeeeeeeeee fflffieess Sesiee/ppape/20//////ee e611cc55778111111333322222233333 Paapee25555--aapP.ppe.. pdf.
Bilal Chughtai, Lawrence Chan, and Neel Nanda. A toy model of universality: Reverse engineering how networks learn group operations. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th In- ternational Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 0243-6267. PMLR, 23-29 Jul 2023. URL https: http:///occeeegss..... press/ sress/v202/ http////////ghtaiaaaa2 wwwgghaitaaaaaaaaaaaaaa html.
Francis Clarke. Nonsmooth analysis in control theory: survey. European Journal of Control, 7 (2-3):145-159, 2001.
Frank H Clarke. Generalized gradients and applications. Transactions of the American Mathemati- cal Society, 205:247-262, 1975.
Alex Damian, Tengyu Ma, and Jason D Lee. Label noise SGD provably prefers flat global minimiz- ers. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 27449-27461. Curran Associates, Inc., 2021. URL https: http::::ss////ooeeiingeeee newr co/ cwpprippp.ccpppprrrrrpppp httpc//pppppaae/e/e//eeeeeeeeeee
Xander Davies, Lauro Langosco, and David Krueger. Unifying grokking and double descent. In www.:neeefffoiiieee
NeurlPS ML Safety Workshop, 2022. URL
tHM
Lijun Ding, Dmitriy Drusvyatskiy, Maryam Fazel, and Zaid Harchaoui. Flat minima generalize for low-rank matrix recovery. arXiv preprint arXiv:2203.03756, 2022.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. volume 97 of Pro- ceedings of Machine Learning Research, pages 1675-1685. PMLR, 09-15 Jun 2019a. URL https: proceed ing ml wtt.sss.rr.....///mmmmmmmmmmmmee v97 du 9c. html.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes peceeparameteiized neural networks. In International Conference Learning Representations, 2019b. URL https: /ttp::///pppeeieeeeeeeee hwt.:/rrt/fffmee/////// id=s1ek3i099
Stanislay Fort and Adam Scherlis. The Goldilocks zone: Towards better understanding of neural network loss landscapes. In Proceedings of the Thirty-Third AAAI Conference on Artificial Intelligence and Thirty-First Innovative Applications of Artificial Intelligence Con- ference and Ninth AAAI Symposium on Educational Advances in Artificial Intelligence, 19A119911119/ AAAI Press, 2019. ISBN 978-1-57735-809-1 doi: 10.1609/aa 609 aaai htt.:/aaa....3333.. 333113573
ai. ai.v33i01 0330131133015744, URL https:/
Andrey Gromov. Grokking modular arithmetic. arXiv preprint arXiv:2301.02677 2023.
Xinran Gu, Kaifeng Lyu, Longbo Huang, and Sanjeev Arora. Why (and when) does local SGD gen- eralize better than SGD? In The Eleventh International Conference on Learning Representations, 2023. URL https:/ htt.httts///ppereiiiiie /tpp//veeeeeeeeeeeee net www.mmirmmeeoooe ttt.:/orrmmmmmssseseeeoem d=sCCouf 6.

---

Published as a conference paper at ICLR 2024
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Sre- bro. Implicit regularization in matrix factorization. In I. Guyon, U. Von Luxburg, S. Ben- gio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https: http//pppr/21///1eeeeee //proceedings neurips wwtpuriips.cppppppppeeeeeepppeee cwtpaper_pppppeeeeeeeeeeeee fi Files/// Sttparrrfssesepaee/////10 1914466aa66671ccddddoooo Papp:e..aaaaaaepppppee. pdf.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th Interaational Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1841. PMLR, 10-15 Jul 2018a. URL htttps://procest ngs. wtal:.ppes//v000
gunasekarl htwprrees//eogguusssaaaaaal 8a htmml....
Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Gamett,, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018b. URL https: http///ooeeeeeeeeeee ngs. neur ps cc/paper_t iles/ hese/aper/2200 8 fi e Oe 98 61 2b 9eb4e4 8a2 2a2 11126688144--2-0- 4c-Paper pdf.
Jeff Z HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Conference on Learning Theory, pages 2315-2357. PMLR, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), December 2015.
Michael Y Hu, Angelica Chen, Naomi Saphra, and Kyunghyun Cho. Latent state models of training dynamics. Transactions on Machine Learning Research, 2023.
Wei Hu, Zhiyuan Li, and Dingli Yu. Simple and effective regularization methods for training on noisily labeled data with generalization guarantee. In International Conference Learning Rep- resentations, 2020. URL https://p/........... http://pppeeeeeeeeeeeeeeeeeeeeee net/ forum?iiiiiooee.ee..... "errm?idiHee33e 3gyHHWH..
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https: http:::://cccccccppeeeeeee www.iinnnnneuuuueee neurips oe paper_ _fl fil pdf.
wttesfilllle es/ http://///....aaaa...aacaaaaaaaa
Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and Lin, editors, Advances in Neural formation Processing Systems, volume 33, pages 17176-17186. Curran Associates, Inc., 2020a. URL https: httttp:://pooceedingg..... newridssnnnn.nnuunoneee.eieiis ccttappe.ccppppe. paiie//paaee///2//EEEae E0le////// TTeeb22tasaF50011111100000011221 wwwpeeeeeb-aaape......eeee pdf.
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve ar- bitrarily small test error with shallow relu networks. In International Conference Learning Representations, 2020b. URL pectps://openreeeee wwwerreeeeeieeeeeeeeeeeeeeeeeeee net foreet/foummmii httpp//prreieeeeeeeeooooooaoaaaa
V. Koltchinskii and D. Panchenko. Empirical Margin Distributions and Bounding the Generalization Error of Combined Classifiers. The Annals of Statistics, 30(1):1 -50, 2002. doi: (0/os///o 15362183.
Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the tran- sition from lazy to rich training dynamics. In The Twelfth International Conference on Learning Representations, 2024. URL https:t///// http://pppeeeieeeeeee //
net/ fenss/ffrree forum? ii
Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetric maximum mar- gin bias of guasi-homogeneous neural networks. In The Eleventh International Conference Learning Representations, 2023. URL https: //ppnreveeeeeeeeeee .ne.ieeww..teeeeeee..e..... net /eet/forrie Forum?ide IM4xp7 5GI
12

---

Published as a conference paper at ICLR 2024
Noam Itzhak Levi, Alon Beck, and Yohai Bar-Sinai. Grokking in linear estimators solvable model that groks without understanding. In The Tweifth International Conference on Learning Representations, 2024. URL https: http::p:://ppereeeee...e.o htt.://eee...oommm..oo/lo11.....
Aitor Lewkowycz and Guy Gur-Ari. On the training dynamics of deep networks with 1.2 regular- ization. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 4790-4799. Curran Associates, Inc., 2020. URL https: http:/::/pppceeeiieeeeee weweeiinnnnnneunseeeeeeeeee.riee nwwrips cc cctpaper.ppppppeeeeeppe lile_fiees//eee wwtps//ppppe///////.///sp pdf.
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations, 2021. URL https: : http://////ee....////m/ net/ forum?id-AHoO Lotp:r/ewweeee..eee.oooo..oaaaaa 7 Sm5H7R.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In International Conference on Learning Representations, 2022. URL
https:://ppenee net
Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. To- wards understanding grokking: An effective theory of representation learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Informa- tion Processing Systems, volume 35, pages 34651-34663, Curran Associates, Inc., 2022. URL https: http://cccccccccceepeepepcecccct newrips cc paper httppcccpppppeeeeeeeeeeeeeeeeeee Satlee//apee/20000 (1t02/2/////df11 0e81 http://weeececcccccacaaaaaaaaaaa
Ziming Liu, Eric Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data.
In The Eleventh International Conference on Learning Representations, 2023. URL https:
/ //pppeeevieww... www.//ewwwwwtmnieeeeeeeeo -tt
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference Learning Representations, 2020. URL /ttttt::://ppeeeee
ew net/ forum? dit.miiddsSSeeeeBBBB..
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of nor- malization layers: Sharpness reduction. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 34689-3644449 34708. Curran Associates, Inc., 2022. /tt.tttppss//proeeeens..ouo s. c.cp/pper fiwes/ wwwes//papeee///////ee pttps//pppppeee////ee/e/ file http:///22....eeeeeeee22222eeee2 /ff.L22255111177777733300002222 3- S-Paper-Confereene .dff.
William Merrill, Nikolaos Tsilivis, and Aman Shukla. tale of two cireuits: Grokking as competi- tion of sparse and dense subnetworks. In ICLR 2023 Workshop Mathematical and Empirical Understanding of Foundation Models, 2023.
Beren Millidge. Grokking. 2022. URL htt.0//tt////....r....////222220 g-Grokking/-
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of Machine Learning. MIT press, 2018.
Edward Moroshko, Blake E Woodworth, Suriya Gunasekar, Jason Lee, Nati Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. In H. Larochelle, M. Ranzato, Hadsell, M.F. Balcan, and Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 22182-22199 Curran Associates, Inc., 2020. www.:i.....uuuure cwtpssscccpppp...pee files wtt.ff/e//pppee/2///iiii... .df.
URL c222ccc9661c170bee77771177666863
Depen Morwani, Benjamin L. Edelman, Costin-Andrei Oncescu, Rosie Zhao, and Sham M. Kakade, Feature emergence via margin maximization: studies in algebraic tasks. In The Twelfth International Conference on Learning Representations, 2024. URL https: https://ogenzewiiee //ppnreview .net// forumm i d=i 9wD
13

---

Published as a conference paper at ICLR 2024
Mor Shpigel Naeson, Suriya Gunasekar, Jason Lee, Nathan Srebro, and Daniel Soudry. Lexico- graphic and depth-sensitive margins in homogeneous and non non-homogencous deep models. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 4683-4692 PMLR, 09-15 Jun 2019a. URL https: httpaa://proceeee /tpp//rceeeddeeeee ngs mol.mmrr..... press/ v97 htt//aacoonn111 httrsonnna..m.....
Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. In Kamalika Chaudhuri and Masashi Sugiyama, editors, Proceedings of the Twenty-Second International Con- ference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pages 3420-3428. PMLR, 16-18 Apr 2019b. URL https: /ptppppeeeengga..p/eepp/// ml r.press/v8 http://vvv///noonn1/ hww./accconnnncooccooooo 9b. html.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learn- ing Representations, 2023. URL ht https: http:p::///peee//eee
oppnreviev. net/ (oorum?id==XFaaccmnme
Pascal Notsawo Jr, Hattie Zhou, Mohammad Pezeshki, Irina Rish, Guillaume Dumas, et al. Predict- ing grokking long before happens: A look into the loss landscape of models which grok. arXiv preprint arXiv:2306.1322 2023.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datasets. arXiv preprint arXlv:2201.02177, 2022.
Adityanarayanan Radhakrishnan, Daniel Beaglehole, Parthe Pandit, and Mikhail Belkin. Mecha- nism of feature learning in deep fully connected networks and kernel machines that recursively learn features. arXiv preprint 2rXiv:2212.
2022.
Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 21174-21187. Curran Associates, Inc., 2020. URL https: / http//ppcceenn..ee/ee 0/file/ (twill//21e2555f8ff
meeeeeddiinnnnnneris..ee...eeees neurips. wttpppsrrcpppppppppppeeeepeee weeeffielslpppee weleeeeapeee//00 pdf.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19 (1):2822-2877 2018
Matus Telgarsky. Feature selection and low test error in shallow low-rotation reLU networks. In The Eleventh International Conference on Learning Representationss 2023. URL https://op enreview net/
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint 2rXiv:2206.048) 2022.
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint arXiv:2309.02390. 2023.
Bohan Wang, Qi Meng, Wei Chen, and Tie-Yan Liu. The implicit bias for adaptive optimiza- tion algorithms on homogeneous neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Pro- ceedings of Machine Learning Research, pages 10849-108554 PMLR, 18-24 Jul 2021. URL mlr wttppp///..//wwggg//// html.
Alistair Watson. Characterization of the subdifferential of some matrix norms. Linear algebra and its applications, 170(0):33-45, 1992.
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How sharpness-and aware minimization minimizes sharp- ness? In The Eleventh International Conference on Learning Representations, 2023. URL https: /tsps://pperreviee net forum?
14

---

Published as a conference paper at ICLR 2024
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635-3673, PMLR, 2020.
Zhiwei Xu, Yutong Wang, Spencer Frei, Gal Vardi, and Wei Hu. Benign overfitting and grokking in reLU networks for XOR cluster data. In The Twelfth International Conference on Learning Representations, 2024. URL https: /ttp///peeeee..eee///ee/e n httt://orimiiieBlllppptt0
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep ReLU networks. Machine learning, 109:467-492, 2020.
Bojan Zunkovič and Enej Ilievski. Grokking phase transitions in learning local rules with gradient descent. arXiv preprint 1rXiv:220011544 2022.
15

---

Published as a conference paper at ICLR 2024
A ADDITIONAL RELATED WORKS
Identifying Mechanisms that Cause Grokking. Some works attempted to identify mechanisms that cause grokking through experiments. Davies et al. (2022) hypothesized that both grokking and double descent are the result of the existence of two patterns: one pattern is faster to learn but gen- eralizes poorly, and the other pattern is slower to learn but generalizes well. However, it is unclear how these notions can be rigorously defined for general training dynamics. Similar to the work of Liu et al. (2023) that we discussed in Section 2, Varma et al. (2023) hypothesized that grokking happens when the generalizable solutions have a smaller parameter norm than the overfitting solu- tions, but the former takes longer time to learn. They also demonstrated that violating a part of this condition indeed changes the generalization behavior from grokking to "ungrokking" and "semi- grokking". two interesting phenomena they named in the paper. For training with adaptive gradient methods, Thilak et al. (2022) identified the Slingshot mechanism, which refers to a cyclic behavior of the last-layer parameter norm that empirically causes grokking. Notsawo cussed similar relationship between the oscillations in training loss and grokking. For the task of learning sparse parity, Merrill et al. (2023) empirically observed that the number of active neurons significantly decrease as the test accuracy starts to improve. While many explanations of grokking may make sense intuitively, none of these works provides rigorous justification their claims with mathematical analysis for neural net training, while our work is grounded by theoretical analyses of implicit biases in the kernel and rich regimes.
al. (2023) also dis-
Progress Measures of Grokking. Efforts have been made to find progress measures that are im- proving smoothly and are predictive of the time to perfect generalization. Nanda et al. (2023); Chughtai et al. (2023); Gromov (2023) focused more general group operations, and attempted to reverse engineer the weights of neural nets af- ter grokking. It was found that the weights can exhibit special structures to make the neural net internally compute trigonometric functions or other representations of group elements. These works also made efforts to define progress measures for grokking based on this understanding of the final weights. In related study, Morwani et al. (2024) derived analytical formulas for the weights of two-layer neural nets, assuming the margin is ultimately maximized. Hu et al. (2023) computed a variety of metrics throughout training and fit Hidden Markov Model (HMM) over them to study the phase transition in grokking. A common issue with these works that why progress measures themselves can make progress is still not well understood and requires new insights into the training dynamics, which is the focus of our work.
the tasks of learning modular arithmetic or
Dynamical Analysis of Grokking. A few recent works were devoted to theoretically analyze the training dynamics in the grokking phenomenon. Liu et al. (2022) attributed the delayed generaliza- tion in grokking the slow learning speed of the input embeddings. They especially analyzed the dynamics of the input embeddings in a related optimization problem, and showed that the eigen- values of the coefficient matrix in the ODE can be used to predict the convergence. However, is unclear how this relates to the training dynamics of the original problem and why the neural net overfits the dataset before the predicted convergence time. Žunkovič and Ilievski (2022) analyzed simple linear classification problem, where the loss is the squared loss with explicit & and & regularization. Assuming specific input distributions, they proved sharp transition in test accuracy and derived an estimate of the grokking time, which depends on the initialization and regularization strength. Levi et al. (2024) leveraged random matrix theory to analyze grokking in a standard linear regression setting, where the test accuracy can provably exhibit sharp transition if the accuracy is defined as the fraction of points whose regression loss is smaller than a small threshold. All these works are limited to linear models, while our work is able to analyze deep homogeneous neural nets.
Concurrent Works. Concurrent to our work, Kumar et al. (2024) hypothesized that grokking is due to a transition from lazy to rich regimes, and observed that manipulating the scale parameter for the output and the task-model alignment between NTK and ground-truth labels can control the time to escape the lazy regime and the test accuracy when staying in the lazy regime. This observation is consistent with our work, but we provide rigorous analyses for both lazy and rich regimes with quantitative bounds. Xu et al. (2024) proved that training two-layer neural nets on XOR cluster data with noisy labels can exhibit grokking. Both their work and our work go beyond linear models, but their analysis does not show whether the transition in test accuracy is sharp or not, while our work provides quantitative bounds for the sharp transition in grokking.
16

---

Published as a conference paper at ICLR 2024
Implicit Bias. A line of works seek to characterize the implicit bias of optimization algorithms that drive them to generalizable solutions. Several forms of implicit bias have been considered, including equivalence to NTK (Du et al., 2019b;a; Allen-Zhu et al., 2019; Zou al., 2020; Chizat al., 2019; Arora et al., 2019b; Ji and Telgarsky, 2020b; Cao and Gu, 2019), margin maximization (Soudry al., 2018; Nacson al., 2019a; Lyu and Li, 2020; Ji and Telgarsky, 2020a), parameter norm minimization (Gunasekar et al., 2017; 201 8a; Arora et al., 2019a) and sharpness minimization (Blanc al., 2020; Damian et al., 2021; HaoChen al., 2021; Li et al., 2022: Lyu et al., 2022; Gu al., 2023). In this work, we characterize the early phase implicit bias based on NTK and late phase implicit bias based on margin maximization.
17

---

Published as a conference paper at ICLR 2024
train
test
100%
800
100%
10
80%
80%
accuracy
60%
10
400
40%
40%
20%
200
20%
10a
0%
25
50
75
100
125
150
100
(3+0)
l0 (1+0
(a) time in log scale (shorter horizon)
(b) time in double log scale (Ionger horizon)
Figure 4: Training two-layer ReLU nets for modular addition exhibits grokking without weight decay, but the transition is no longer sharp: it can take exponentially longer than the time needed for reaching 100% training accuracy.
B DISCUSSION: OTHER IMPLICIT BIASES CAN ALSO INDUCE GROKKING
The main paper focuses the grokking phenomenon induced by large initialization and small weight decay, where the transition the implicit bias is very sharp: - log a leads to the kernel predictor, but increasing it slightly to L+s log a leads to a KKT solution of sunn-norm/max-mmrraee problems associated with the neural net. In this section, we discuss grokking: implicit margin maximization and sharpness reduction. These two sources can also cause the neural net to generalize eventually, but the transition in test accuracy is not as sharp as the case of using large initialization and small weight decay.
-
other possible sources of
B.1 IMPLICIT MARGIN MAXIMIZATION WITHOUU WEIGITT DECAY
One may wonder if grokking still exists if we remove the small weight decay our setting. In Liu et al. (2023), it has been argued that adding non-zero weight decay is crucial for grokking since the weight norm has to lie in a certain range (Goldilocks zone) to generalize.
We note that for classification tasks, Lyu and Li (2020); Ji and Telgarsky (2020a) have shown that gradient descent/flow on homogeneous neural nets converges to KKT points of the margin maxi- mization problem (R1), even without weight decay. So it is still possible that some implicit bias traps the solution in early phase, such as the implicit bias to NTK Allen-Zhu et al. (2019); ?); Ji and Telgarsky (2020b); Cao and Gu (2019). Then after training for sufficiently long time, the margin maximization bias arises and leads to grokking.
We empirically show that training two-layer neural nets without weight decay indeed leads to grokking, but the transition is no longer sharp. More specifically, we use the same setting as in Sec- tion 2 but with weight decay À 0. We focus on analyzing the training dynamics of full-batch gradient descent with learning rate so small that the trajectory is close to the gradient flow, which ensures that gradient noise/training instability is not potential source of implicit regularization.
However, directly training with a constant learning rate leads to very slow loss convergence in the late phase. Inspired by Lyu and Li (2020); Nacson et al. (2019b), we set the learning rate as Z(0())) where 7/ is a constant. The main intuition behind this is that the Hessian of L(0) can be shown to be bounded by ((C(0) poly([ (|]]))))) and gradient descent stays close to gradient flow when the learning rate is much lower than the reciprocal of the smoothness. Figure 4 plots training and test curves, where z-axis shows the corresponding continuous time calculated by summing the learning rates up to this point. It is shown that very quick convergence in training accuracy occurs in the early phase, but the test accuracy improves extremely slowly in the late phase: the test accuracy slowly transitions from ~ 0% to 100% as t increases from 10 to 10 In contrast, Figure shows that training with weight decay leads to a sharp transition when the number of steps is ~ 10?.
18

---

Published as a conference paper at ICLR 2024
We note that this much longer time for grokking is not surprising, since the convergence max- margin solutions is already as slow as O lng logt for linear logistic regression (Soudry al., 2018).
log
B.2 SHARPNESS REDUCTION WITH LABEL NOISE
Grokking can also be triggered by the sharpness-reduction implicit bias of optimization algorithm, e.g., label noise SGD (Blanc et al., 2020; Damian et al., 2021; Li et al., 2022) and I-SAM (Wen et al., 2023). The high-level idea here is that the minimizers of the training loss connect as Riemannian manifold due to the overparametrization in the model. Training usually contains two phases for algorithms which implicitly minimize the sharpness when the learning rate is small, where in first phase the algorithm minimizes the training loss and finds minimizer with poor generalization, while in the second phase the algorithm implicitly minimizes the sharpness along the manifold of the minimizers, and the model groks when the sharpness is sufficiently small. The second phase is typically longer than the first phase by magnitude and this causes grokking.
As concrete example, Li al. (2022) shows that grokking happens when training two-layer di- agonal linear networks (4) by label noise SGD (6) when the learning rate is sufficiently small. We consider the same initialization in Section 3.2.3, where (()) a(1, 1), where
0 controls the initialization scale. Different from Section 3.2.3, we consider & loss, and the label is generated by sparse linear ground-truth, that is, y (æ) x: "w", where w² is k-sparse. We also assume data distribution is the isotropic Gaussian distribution (000 f).
O + 1) = ((o()))) - ys. + E,)?, where & i.d. N(O,
(6)
Theorem B.1 There exists C₁ > 0, C2 E (0, 1), such that for any n, d satisfying Cislud S n S
d(1 - C2), the following holds with Qin)) probability with respect to randomness of training data:
1. The gradient flow starting of C from (()) has limit, OGF. which achieves zero training loss. The expected population loss OGF is at least [ (Theorem 6.7)
2. For any a E (1,2) and T > O, (TT/7²) converges to Per in distribution as 1 - 0. (Implica- tion of Theorem B.9)
3. For any T 0, O(T (1²) converges some ( in distribution as + and ( has zero training loss. Moreover, the population loss of (TT) converges to as T + co. (Theorem 6.1)
The above theorem shows that a sharp transition of the implicit biases can be seen in the log time scale: when log log log for any a (1, 2), it converges dGr; but increasing the time slightly in the log scale to log log T + 2 2og 1 leads to ((T),, which minimizes the population loss O. However, this transition is less sharp than our main results, where show that increasing the time slightly in the normal scale, from A
1 log a to X
1+c log a, changes the implicit bias.
19

---

Published as a conference paper at ICLR 2024
C PROOFS FOR CLASSIFICATION WITH L2 REGULARIZATION
C.I PRELIMINARIES
Let FiL [1, +00) + 00,) + x log x be the linear-times-log function, and FiL. : [0, +00+ + (oc+ be its inverse function.
Lemma C.1 For all
Fi 2 log(s+1)
Proof: Since FLL an increasing function, it suffices to show FiL logloty)) S y, and this can be indeed proved by
log log[v+I]
FLL
1) - log(y 1)
y log log(y + +1 log (y 1)
S y,
log(y
where the last inequality holds because log(u+1) S
Lemma C.2 For all and 3 E (0,
Fi. - 1 2 ((F 1).
By concavity of FIL FIL (By) 2 BFT (v) + (1 -B) gives the desired result.
Proof:
Rearranging the terms
Lemma C.3 The following holds for "1....,fr In and L = Σ r2
1 - < L log 1 L
C Cogg
ri log
nL
Proof: For proving the first inequality, note that log r 2 log - 5o
ne.
n
log
log nf r;
log
Ti
nL
For the other inequality, by Jensen's 's inequality, we have
1 Fa)()s) S- (C) = Llog
ry log r3
L
which completes the proof.
C.2 KERNEL REGIME
Now present the proof of Theorem 3.4 for the kernel regime. We follow the notations in Sec- tion 3.2.1 and define h hit as the unique unit vector along the direction of the optimal solution to (K1) and Yntk := min [y (Vfi (00itt, h. The kernel regime a training regime where the pa- rameter does not move very far from the initialization. Mathematically, let max constant so that
0 be small
Yutk
max - : i [n]} < 2
sEin]
holds for all ||0 - 2 < maxx and let Tmax == inf {1 f12 1200 0: ea 0(1) O
> 6max
j be the
time that (()) moves far enough to change the NTK features significantly.
A key difference to existing analyses of the kernel regime is that our setting has weight decay, so the norm can change a lot even when the parameter direction has not changed much yet. Therefore, we normalize the parameter by multiplying a factor of Att and compare it to the initial parameter direc- tion. We consider (t) as moving too far only if || | ((t) m l2 too large, which different from existing analyses in that they usually consider the absolute moving distance ||@(t) ((0)
20

---

Published as a conference paper at ICLR 2024
In the following, we first derive upper bound for the loss convergence rate, and then use it to lower bound the training time in the kernel regime Tmax 2 3. (log a log log A + 9(1)). Finally, we establish the implicit bias towards the kernel predictor with more careful analysis of the trajectory.
More notations are needed for the proof. Let (;( == (0) be the output margin of the neural net the i-th training sample. Let (0) ((). U/) exp( (00(--;(()) (0)) be the loss incurred at the i-th training sample. Let A = 22 which is factor that will appear very often in our proof. Let Gntk sup V f: [|2 ||0 mitt 2|2 < Cmax, i E [n]]} be the maximum norm of NTK features. The following lemma gives an upper bound for the training loss over time.
Lemma C.4 For all S S T max.
dL
dt v a 2((6-1))- L² + AL Llog 1
litk
L'
22(2-11At
((O(1)) S max
(8)
1
Yonk I6[L-1]
A(l - ( Pi..
Yoak
A)
8L
Proof: By the update rule of 0, the following holds for all 0:
dc dt
do de
AL
=
T4
T1
=
=:V2
For V1, we have
V1
=
ri
V. 0
2
ri (V); Mick)
(9)
n ==1
Note that for all t < Tmaxx
Ti (0)()), )
1
Σ
ri (VVg((ae Quutt)) Whek) - (() fila -
n
1
T a L 1(--)11 (Nuek 2 Yatk
1
1=1
Yntk
1-1 (1-1)
a
2
where the second inequality is due to (Vq: hanx) 2 Yutk. f 0(t)) f mm)m))))
<
Yatk and the (L 1)-homogeneity of fv. Together with (9), we obtain the following upper bound
2
for Vi:
2 2(1
Vi /
nik
a
C
(10)
4
For Vz, by Lemma C.3 we have En ri log 1 + S Clog 1. So we have the following upper bound for Vj:
V/2 s XL Clog -
(11)
Combining this with (10) together proves (7).
For proving (8), first divide by L on both sides of (7) get
d 1 1 2(L-1)
>
XL
log
(12)
d Σ 4
L
21

---

Published as a conference paper at ICLR 2024
Let & == {1 (0, T max] F()(2))
Iup 82
Ae AC-2(L-120 For all [0, T maxx], if E, then
by Lemma C.2,
1
C(0(1)) -
FILL
Li 87.
AE -2-2(12 (
1 1+
FiLl
(
Link A) - 1) 2((L-1) 8I.
1
c2(4-1)At
FIL
Yas 81.
A) 1) 2(1--)AA
FLT.
Titk 8L
A)
Otherwise, let l' be the largest number in U {0) that is smaller than Then
1
d
+
dr
((O())) ((0(1')
dr L
2 C(O(V))
1
+
["
Yutk
a (LL-1) -2(C-11) dr
1r
8
2
C(O(t')) + 16(L
Talk
Ale -((-2111) - - -2((-))))
If E' E E, then Z[e(e] 2 PT..
Trnk 8L
Ae 2(L 1)Ae") 2(1
2 Pi..
Trsk
Ae -11L 1)At) 1/A1). -((--))A
which contradicts
SL
to t & E. So it must hold that t' = 0, and thus Z((())
A(1 e 2(- 1(A))))r). Combining 2 1+ 18(1-1)
Tull
all these together proves (8).
Lemma C.5 There exists a constant AT such that Tmax 2 1 (log 1 log log A + AT) when X is sufficiently large. Furthermore, for all t S } (loga 1, log log A ++ AT).
1 MO(() O
Oa log A) etil
Proof: By definition of T maxx prove the first claim, suffices show that there exists constant AT such that ! 0t) Omit|| l|2 max for all i S min (log a T log log A + AT) T Tmax):
When < Tiax, we have the following gradient upper bound:
|VC(1())
(o(())
(o((((
1=1
n:a - (L- 11-1) D Va
1=1
< Gauaka (1-1) C.
By update rule, ett(()) a0 = 1/' e VL((( Applying the gradient upper bounds gives
1 Me))) aminn 112 < / ||>))(((())]]) dr
< /--
((((T)) dr
Guiatt2 (' (0-(222) ((0(T)) dr.
Substituting in for ((L-))) ((((T)) dr with the loss upper bound in Lemma C.4, we obtain the
following for all 1. S Tiaxx
(L 2-
e ((L-2) 2)AT 1((4-1)A7
(L- 2) ((0(T))
max
dr.
Yonk
A(1 (((--)
FrJ. 8L
2 A)
16
22

---

Published as a conference paper at ICLR 2024
So have Jo e- (1--2)A7 (((2))))(((((T) dr I max(41, 12), where
dr
ELAT
dr.
1+
Tirk A(1 - 1(1m)))7)) 16(L-1)
FIL
Tack 8L
A)
-
We can exactly compute the integrals for both and 12. For It, it holds for S min O( log
loga)) } that
log (1+ 11
=
16(L 1) 1
Yatk
22(4-11Ar
16(L 1) A
Yakk
2(L 111-111( (1 IG(L-1j A
log (1+
2 16(L - 1) A poly(a)
Yutk
2(L 111 (1
You'
A
16(L - 1)
16
=0 =0
log A
(AA XA
For I2, it holds for all t Tmax that
e LL
log A AA
1 = LA
= =0
P
A) -1
8L
Putting all these together, we have
1
11 0(aa(a1)) (o
log A
log A
0())((1) aliin
+ 0
AA
AA
(13)
0 a a A
cliee 0(a 'log A) etiee
AA
Therefore, there exists a constant C > such that when a is sufficiently large, ((t) Qmittly I Ca- log A elir for all t S min(Tmmax, O log Setting AT = log Eamas we obtain the following bound for all min(Tmax; (log + log log AT)):
00()) Omit
S Ca log exp (L a log log A log
Cmax
Emax:
which implies Tmax } (loga claim.
1 log log A AT). Combining this with (13) proves the second
Now turn to analyze the implicit bias. Let 6(t; a0mit)
(2 e admit)
O and
htt: aQuit) = a-ee- -LM 8(t a0uit). The following lemma shows that the output function can be approximated by linear function.
Lemma C.6 For all S A (log a log log A AT).
()) (f 8) O(a -" log A) c/A0|10112
(14)
(0) (O h) + -0 log A) cL |
(15)
Proof: By Lemma C.5 and Taylor expansion, we have
a((e"00) 0) q: (iiit) + (Onit) 8) 0(1811
= q: (mitt + V q: (iit) 8) + O(a -L log A) LAC 11
(O 8) + O 00a log A)
cLAi ||1||
which proves (14). Combining this with the L-homogeneity of f: proves (15).
the following two lemmas, we derive lower bounds for the loss convergence that will used in the implicit bias analysis later.
23

---

Published as a conference paper at ICLR 2024
Lemma C.7 For all 0 < T max.
dC
dt - 2 -
2(1.-1)At AI. Llog 1 nL
(16)
((0(t)) 2 min
FiL (22 Gr Ae 2(L-)))))) 'n exp(1 2(L ((L-1))A)
(17)
G
The proof is similar to that of Lemma C.4, but now we are proving lower bounds. By the update rule of 0, the following holds for all > 0:
Proof:
dC LC)))
d0
AL
q
Ti log
n
r3s
=:V4
1/2
For V1, can lower bound it as follows:
1
Vi 2
r3 2i (0)
n
211 2(1-11 2(1-1)1
Vi L (0))
11
- 2 l) 2( 1)M
C2
where the last inequality holds because ( 0)1l2 (0))l1 S [Ven(Omit) + ( Vq, (0m)))2) S Gatk-
For 1/2 by Lemma C.3 we have 1 E log 1 > L log म So V/2 2 AL Clog z Combining
r;
the inequalities for V1 and V2 together proves (16).
For proving (8), we first divide by L on both sides of (16) to get
dt C S G Gma
2l--1) 2(1-1)A 2(2-11- - AL log
1 1
(18)
L nL
Let E = E 10, Tmax] : ( mz) <
nt. ntk AE 2(L-11 It is easy to see that 0 E E, For all t [0, Tmax], if E E, then by Lemma C.2,
2 G2
((O(1))
n
FIL mL
G2 ntk Ae Ac-2(1-11))
no t at the boundary of E, then we are done. Otherwise, let t > 0 be one on the boundary of Then at time
d 1
AL 1
log
dt mL
2
nL
So as long as 1 log (C()())
2(L 1)A,
Fi,. 1
d
= (1 log
di nL
nL
dl. nL
1
XL
1
--1 + log
log
nL
n.
nC
-1 11 + log
1
G2
Cak
Ae 2(L-1)At
(
(2
nT. mikk AE- A 120-
mL
nL
which means / is not actually on the boundary. So 1 + log neterey S 2(L - 1)A. Repeating this
argument with inequalities and taking some integrals can show that every point outside & satisfies
I log n(((()) S 2(L 1)A
24

---

Published as a conference paper at ICLR 2024
Lemma C.8 For any constant E (0, letting Te (a) x
log a, it holds that
d 2lll a 2(11111) 2((--))
h
(H)
L
(19)
dt
|
1 log
(Tatk + O(a log A) elary dl
(20)
dt
C
dt
Proof: By update rule,
dh
= e
LAd dd
L
dt
dt
e LLM
(((()) LAh
i=1
(11-11))~2111
2((1-1)AM
r (01)) LAh
Then we can deduce (19) from 4 |///2 = ar
dh
T
For proving (20), first, we apply the chain rule to get
d 1
X
dt L [()) (- (-V(()) 0o, VL()))
log
(()) (0, (0-))
((
1
For V1, note that
1 n
1 n
|V((0)
r: (
L- -LL-1)A
ri(0) (
n
=:Vs
By Lemma C.5, [|{(t]]2 O(aa log A) e (ze
(Omit)) h h +
O log A) eLde One way give lower bound for V3 to consider the projection of 1 - (0) (2e20 along hi
V/3 2
(H M) hick
Σ (0) ((Vg (0 hick) + O(a= A) (d0))
2 (()) (Tmek + O(a log A) elwe) > 0.
Another way to give a lower bound for V₂ to consider the projection of E (H along h:
V/2 2
(H
A
TA
Putting these two lower bounds together gives:
V2
Vi = ((0)
1
2 (Nack + OO(a-"log A) E rd a
r:
n
1=1
25

---

Published as a conference paper at ICLR 2024
For V2, by Lemma C.6, have
X
V2
(0) (0. Vq.(0))
( n
LA
(0)
nC(0)
LA nC(O)
(0) ((Vg:(0mit)), O(a A) (e)
< LAllhll2(Titts + O(a log A) e eLAc).
Putting the bounds for V1 and 1/2 together proves (20).
Now we are ready to prove the directional convergence h(T (a:); mm to hake
Theorem C.9 For any constant e E (0, 1), letting TT
(a) = - log a, it holds that
hTT (a); an
lim a-+
hitk.
log X
Task
Proof: Integrating (20) in Lemma C.8 over / E [0,To (a)) gives
log ((O (a))) (Titk O(a log A)) [1h(T, (a))
By (8) in Lemma C.4 and (17) in Lemma C.7, we have
log L(O(TE (a)))
log(417. 1) - 2(L -1) (a)
log A - 2(L - 1) - c) log + O(log log A)
log x O(log log A) +oo.
By Lemma C.6, log D(0(TE (a)))
< (()) so we also have [| (a))lla - to0.
Then
minss[]]]
(( h)
mingg(n) (4))))
Yntk 2
+ Oaa llo A) -e e2TT (a)
|h/l2
|///2
log ( En exp(-q: + Oaa log A)
|
log 2(0) Ola-ol log A) Yutk.
h2ll2
Therefore, we have minge(x] V (Oiitt)
ICT, (a)):ammm)) [[A(T] (a)
)
- Yatk- By the uniqueness of the
A(2 (a):a0mnic)
max-margin solution of (KI), it must hold that
hith as + +00.
[[a[T, (a)aamal)l
Then, we can approximate the norm of h as log
log a + O(log A)
log a (1 +1+00 o11)
[h l|22
Tntk + o(1)
Yutk + o(1)
Yotk X
So h can be approximated by
h
1
log a hitk
h =
|/1l22 |
(1 o(1))
7ntk
X
which completes the proof.
26

---

Published as a conference paper at ICLR 2024
Now are ready to prove Theorem 3.4.
Proof: For any imput æ, by a simple extension of Lemma C.6 to x+ we have
f(0: x) (V x), hi) + O(( "loo log A) _ 11ell2.
Normalizing by (a) and taking limits, we have
[(0:z)
( x), Z(a) + +Oaa Oaaattaaaa -lo A) e LMM ||A||2
+ (VS(Omiit; x), hitk))
((a)
Z(a:)
which follows from Theorem C.9.
C.3
RICH REGIME
Now we proceed to prove Theorem 3.5 for the rich regime. First, we derive the following bound for the training loss at the end of the kernel regime. Lemma
C.10 Ar t = 3 (log a log log A + AT), we have Lx(0(t) = (A )/E))
Proof: By Lemma C.4,
easily obtain the bound
L(((L)) = O A exp(2(I. () I log log A))) = O'Allg 1)-1+2/1.)
Also we have the norm bound ||0(t]l2 = Olae >1)
Olog A) A) since is in the kernel regime. We can conclude the proof by noting that Lx + ||0||3 and log O(log })
The loss bound then implies the following norm bound.
Lemma C.11 If (0(0)) = OAA(log ) then for all 2 0,
[[0(t)l22 @(log 11.).
Proof: For all t 2 0, the following relationship can be deduced from the monotonicity of Ex. r.(O(t)) < (O))t, and the Lipschitzness of fr:
1
2(log(1/A)) & log nCx()(0) S log nLx(t(t)) S {:(((()) s B
So 10(t)ll2 = 2((og "))) Also note [[0(t) < (0(t)) s (0(()()) = O(log 11)) So we have the tight norm bound [[O(t]]22 ((log) 114
(((log))11)
To prove Theorem 3.5, suffices to prove the following.
Theorem C.12 For any starting point (a(a)) satisfying L (0o(a)) = O(A(log 2)2:2) and time function T(a)= 2( A log }),, given any sequence {ax} with ak + +oc, there exists a sequence 0(tr:0o(az))
(teh22 such that t I T(a), and every limit point of { (Tot4:0o(ax))2 direction of KKT point of (R1).
2 0} is along the
Proof: For all a > 0, S (x((((a))) = (x(0)) - So T(n) | (O)) Since (Oo(a)) O(Alog ))2/L), we have the bound fo T(a) (Z)((()) dt < O(Aloo 4)2/2), which implies that there exists a time such that [VCx < (1og )) (1-2/11)
Now 4 = Tax suffices to prove in the case where ((r)))((aa)) E S converges; otherwise (o 2)
we can do the same prove for any convergent subsequence. Denote 0(t;: (o(a))) as ((k) and ((x)))
(t(t::0laa )112 as g()) for short. We claim that the limit 0 == lim ((k) is along KKT-margin direction of (R1).
First, according our choice the we have
( (a(()))) - (l(k)
< O((llg A) -1111/21//D)
(21)
t=1
27

---

Published as a conference paper at ICLR 2024
By Lemma C,11, 2|2 O(A(log 11). Now divide x||0(k)|22 both sides of (21):
|0112-1 1-2 ((())) (a(()))) -
o (log A) ()-(1/2-1/1) (log 211/2.
(22)
nd
|10 4-2 ())))))) We can show that μ4,s O(1) by
Let μ4 be vector with coordinates Mk,
noticing /: (((())) T0ll ((((())) = (1) and
Mk Mk. ((((e) S
Mk (i
+
Mk.iy V (i: -0()) 0(6)
=1+0
(log A) (11221111//)) 0(1).
log X
Then, let {(ax, tx2 1p21 be subsequence of {(au, tx ) J221 1 that MR, R" converges. Let á := -imo-++00 fkp, be the corresponding limit. We can take limit hp + too on both sides of (22)
and obtain Σ Pays V ((()) = 0.
Let i, E [n]] be an index such that qs (0) Jmin (0). We can verify that μ. for all E [n] with f() > qmin(() by noting fu exp(1,, (0 - fo (0(*)) = exp 0 (q. (0) (0) +
and f + 0.
Finally, Theorem 3.5 can be proved by simply combining Lemma C.10 with Theorem C.12.
28

---

Published as a conference paper at ICLR 2024
D PROOFS FOR REGRESSION WITH L2 REGULARIZATION
The analyses for the regression task are similar to those for the classification task, but the use of squared loss simplifies the analysis.
We define 8; () : -f(0)) y:- In the regression setting, the gradients of the vanilla and regularized square loss can be written as
VL(0)= = 2 Σ s; (.).
(23)
i=
and
VZ2 =2 Σ (.) +
(24)
D.1 KERNEL REGIME
Now present the proof of Theorem 3.9 for the kernel regime. Following Section 3.3.1, we define hit as the solution to (K2). Additionally, let (0) E Rxxn be the matrix where the i-th column is (O). Let (()) = (0) (0) be the NTK matrix at 0. Let Pntk > 0 be the minimum eigenvalue of K (0iit) Let Ema be small constant so that Amin (nin(K(0))
Mutk holds for all
2
||0 iiitlle la Cmax- Let T max == inf 0: ((t) miit 2
> 6max }. which the time that
(t) moves far enough to change the NTK features significantly.
First, we derive an upper bound for the loss convergence.
Lemma D.1 For all 0 S t S. maxx.
dC lotk a (L--10 22(- (
AL
dt <
L+ Vn
llull2vE,
(25)
2n
[(((()) & max
{ { Ill1l12 exp
(- V
18([L-1) (1 (1-
(1=
22 a ((L-))
16nL211e113x 4( 11Ae
c
(26)
Proof: By the update rule of θ, the following holds for all t 0:
2
dC
O d0
12
XL n2
8
f
stfe
71
=:Vi
For Vi, we have
11=
r
V (0)
s0) T (())((())
n2
2(L 1) 2(L- s(0) K( a ()s
n2
2
Mntk
a
(()-1) -2(L -((L-11 (s))]11
2n²
Dutk a
2(1-11)
1) 2(LL-1)
2n
For V2, we have 1 -i A + s.y,) 2 + Σー st 2/- -
So 1/2 =
XL Vlyle Z. Combining the inequalities for Vi and V22 together proves (25). v
For proving (26), we first divide by L on both sides of (25) to get
d
log L <
Putk a (L--1)
2(1 1)
AL Vn ||||
1
dt
2n
VC
29

---

Published as a conference paper at ICLR 2024
Let 8 = {t [0, Tmax]:
AL vn
llulll2
v2 2 Yatk ( -2(1-1))) For all t [0, Tmax], if E,
then
2² S
1
L
11
n
(Mick a 1)e (- 2(1-1)2002 in
16n
X²
alL-11) cl(L-1)Ar
ntk
Otherwise, let t' be the largest number in & U {0} that is smaller than t. Then
log ((((t)) log ((0(t)) [ dr
d log ((((T)) dr
f: (Munt 2(L-1)
S log ((((t))
-- 1)rr) dr
4n
log L(O(('))
If E E, then ((O())) < L((())) v 16mL 2lvll
12
ALL 11Ae
16nL"Ilvll Aa
a01L-11 e 4(2-1)A
a012 17
e
Misk
Otherwise, and we have
(-f,'
[()(()) < ((0(0)) exp
(nik (2(11-1) -22(0-111))) dr
4n
PuckA 8n(L - (1-
1 exp
7
which coneludes the proof.
Now we derive lower bound for the time that the dynamics stay in the kernel regime.
Lemma D.2 There exists a constant AT such that Tmax 1 (log AT) when sufficiently large. Furthermore, for all S X (logas a + AT),
| M 0(1)
Onnt = oaa eLee
Proof:
By definition of Taxx, to prove the first claim, it suffices to show that there exists a constant AT such that ((())(()))) Qmit llz S mmax for all /. S min & (log AT), max) )
When S Taxx. we have the following gradient upper bound:
1 n St
VC(0())
vV(())
71
((()) || (0))l2
n
aa 1-1 1o (4 (4-1)) 10 - 0)
(0)
Vn
I a 1-1 -
By update rule, et e*t(t) - a0 = 1ee e (((((T))) Applying the gradient upper bound gives
0mit 112 < 1 /
| ((t)
|
dr
00
< a / - (1 (C(((T))
G fe ((10))) dr.
30

---

Published as a conference paper at ICLR 2024
Now substitute L in So e-(1--)AT [(O(T)) ( dr with the loss upper bound in Lemma D.1. Then for all S maxx. we have 50' - (C(((T)) dr I 1 12), where
(- Vntke (1-1))1111 (1---211-1A) -(4 (2-r))r).
11
11l12 exp
166
vn
4vnLly
A LAT dr.
22(1-1)
Vutk
We now bound 1 and 12. For T, we divide the integral into two parts. Let to = 20A2-1)A- for
T E [0, fol, we have - - -2(2 2 2(L - 1)AT (L - 1)AT, so that
exp
VntkA
- (1---2((-1)rr)
(L (L-2),4 2)
16n
1)
PnttkAAT
-0 (
exp
dr 0
16n
AX
On the other hand, we have
I'- exp (T66((L 166((L-1) 1I
VntkAA
(1- -(L 2)ATT)
< exp (-((()) S." exp (--()(-2)A7)d
-0(A-)=- =0
-0-10
do eexp(-((A))
Q2(L-11
Thus, have
(
11 = Q2(L-1)
For 12, it holds for all & S T Taax that
12 0 (2(11-13 ebi
Putting all these together, we have
l - ab aliittll =0(al
Q2(6-11) elde
(27)
= = elide
Therefore, there exists constant C ( 0 mmitlly I Ca Le LAt for all S min O log a)}. Setting AT := 1. log Smar, obtain the following bound for all t min(Tmax. { a + AT)):
such that when a is sufficiently large,
L Mo(t) Omit 2
Ch exp (Lloga a log
Emax
Cmaxs
C
which implies Tiax 2 + (log + AT). Combining this with (13) prove the second claim.
Let 6(t: ad following, we show via a series of lemmas that for any o that is in the orthogonal complement of Φ (6.z) 0, which is crucial property for deriving the implicit bias.
(1 ( ammit) Omiss and h(t; cmmi)
LL almit). In the
1812
Lemma D.3 For any vector that lies in the orthogonal space of the column space of T. have (8(T) (a); a0 v) = O (a (a-2)) where T (a) = -- = log a.
Proof: By update rule,
d5 dl
1e
-2
8,
fi ((itt))))))))
31

---

Published as a conference paper at ICLR 2024
Then have d v6
2--
(
di
a
1
-2-e a
86 (t( V fi(aee (00(ac-)mmii))) MOmii)),
i=1
22a 2-2 2AA (14-))
admitt f:
where the second equation is due to our choice of and the last equation follows from the homogene- ity of ft. Now, since fi is locally smooth in neighbourhood of Innt and || e m (a--4)- edLi by Lemma D.2, we have
IV - a0mit) - V (i - 0(aa" L LAr)
Let to 20(2-1)X 1)x
then for E [0, tu]. we have (2((1-1) 2 (L 1)Ar, so that
d (6, v)
0
2 2,2 V (0(1;a0intt)
dt
a 2622 max { exp
Vntk AAL
+ St ,0 , -2(1-1) 1) 2(1 })
0
16n
0 max -2 eexp
da 22 22x2
exp
32n
for sufficiently large a, since A in (26) dominates the first term, so that
2(2-1)
On the other hand, for t [to, 2- (a)), the second term
d d v) dt
0
2 22
/ (((t:
O
2222 221-11 (e O(Aa 22 -A 22M))
Hence,
(a)
Pnt
dl, do -26 f (a)
max
a
exp
32n
J
a 2 1
}) = 0 2L 2011-0)
max
AA
which proves the claim.
Lemma D.4 For O(t) with t < 1 (log a + AT) as defined in Lemma D.2, we have
f( 0xx) (V mii x), 8) + O(a 1 f(0:x) (V f0 miti x), h)) + 0(a-4)- e |
(28)
(29)
Proof: By Lemma D.2 and Taylor expansion, we have
0(;) 0xx) f(0mit: x) + (V f miti x). 8) + ((188
5(0mit x) V f 0 miti x
) 6) + O(a / L |0||2
f x), 8) + O(a h) JAAt ||81|2.
which proves (28). Combining this with the L-homogeneity of di proves (29).
Lemma D.5 [h<TT (a:): a0 m2) O(1).
32

---

Published as a conference paper at ICLR 2024
Proof: By Lemma D.2 have (a): amm)) 122 [[ (a): eake-LATE ||
that
O(1.
On the other hand, recall that T, (a) = 1-c
x log a, so by (26) we have
| S VnC(O(T (a); = O max {e 2(1), dAx 2 )1)))a)) = ((-22))1))))
Thus, for sufficiently large a, we have (a): fy.. By Lemma D.4, we have that
yi - |/1 = ||h||2 = 22(1)
0(a
desired.
Corollary D.6 lim
( [18|a (To (a); (a); = O.
By the previous lemma, we have [|ō(T; (a); aniiii))ll = Q a0aaaaa aaa-LelxT (la On the other hand, (8(TZ (a); v) O (a 222) by Lemma D.3. The conclusion immedi- ately follows.
Proof:
Theorem D.7
For any constant E (0, 1), letting To (a) == - log o, it holds that f(O(T, (a); al mit): (V æ)), h
Va Rd
lim a0++00
Proof:
By Lemmas D.4 and D.5 we have ((O(T, (a): mitt);
(Oii( x) h0 (0mit)) 0(a .
(30)
(V (0 a0mit)) ))) (la-ch)
By Lemma D.I we have for Vi [n]. (0 aQ): - -yil O(a 10(a-2((-11)) -((1-1)). Since yı (Oiitt)))) m we have
lim
(a): a0): - h = 0.
(31)
By Corollary D.6 and Lemma D.5 have that for any direction v orthogonal to the column space of 4, we have 1im (h(0(T= (a); a0mi)) v) 0. Let Ps be the projection operator onto the column space of then
lim (I - P4)1(O(TT (0): = 0.
(32)
Combined with (31), we deduce that
(Vii( Poh(O(TT a0iit))))))) - hjikk) = 0. +00
lim
The above holds for all V. hit also lies in the column space of d, so we actually have lima-++0o (P aa):a) haulll 0. Hence, we can use (32) to deduce that lim [I (a): al) h h 0. Finally, plugging this into (30) gives the desired result.
[d]. which are exactly the columns of t. Note also that
D.2 RICH REGIME
Now we proceed to prove Theorem 3.10 for the rich regime. First, we derive norm bound.
Lemma D.8 For any constant > 0, ler T (a) = 1tc log then we have that
X
0 (1; (1iii))))) 1|2 0(1)
max (og
33

---

Published as a conference paper at ICLR 2024
Proof: We know from Lemmas D.1 and D.2 that
L log or; mnmi ) max
Ilyl2 I1vl2exp exp (- 16m(L. - 1) ) 16n 1 0(22).
MntkA
A
Moreover, by Lemma D.2 we know that |10 (1 log or almitt
1|2 =O(1), so that
Lx @0 ({> log a: admit) =c log a; aQmis a ))
\ 0 ( log az; a )1² O(1).
+ 2
Since 10 (1; ((iit)) J:20 is the trajectory of GF on Lx, we know that for 41 1 log a, we have Lx (0 (t; (nitt)) (()) as well. This implies that ||0 (timt)) ||2 ((1) desired.
Now we prove Theorem 3.10.
Proof: By Lemma C.5, know that for t 1 log a, have (10 aQmit) Omitll2 ((1) [10(t: m0 (1). Since
Li (o ( log logg;; al )) - L2 (T (ai); (():)))))))
TT (a)
log
and 0 S La ( that there exists ta O(A(log ()
Lx (0 (1 log a; aDmit)) s O(A) by Lemma D.1, we deduce [ log a, Txt (a)) such that [[Cx(((;;00m))))) s OVV (VAgloga) loga) =
For any sequence {ar)ee with ag +o, we choose
Let {f(ti, ai, m h }} be any convergent subsequence of {0(ts; k: 2 1}, then have that 10m |VL, (o(ti::admi)) 1|2 0, i.e.
law
lim ++0
2
8 VV an Omit)) + AO(1i,) : as Omit)
= 0.
(33)
By the choice of (ix)A2) we know that {((tic au mmit) } JR21 converges a point 0* Rd so the above implies that
Vf((0') X0" =0,
thus 0* is KKT point of (R2), as desired.
34

---

Published as a conference paper at ICLR 2024
E GENERALIZATION ANALYSES FOR LINEAR CLASSIFICATION
In this section, we provide supplementary generalization analyses for the linear classification ex- periments in Section 3.2.3. We follow the notations in Section 3.1. Let the input distribution Dx. and the ground-truth label (x) be sign(æ w²) for an unknown w E RR For training set (Z1,..., define yi = y" (xx) for all E [n]]
E.1 PRELIMINARIES
We base our generalization analyses on Rademacher complexity.
Definition E.1 Given family of functions F mapping from I&d to R, the empirical Rademacher complexity of F for a set of samples S = (2.,...., Fn) is defined as
RsFF) B Bo mmif((
sup fer
off(zz)
(34)
Rademacher complexity useful in deriving margin-based generalization bounds for classification. The following bound is standard in the literature Koltchinskii and Panchenko (2002); Mohri et al. (2018).
Theorem E.2 (Theorem 5.8, Mohri et al. (2018)) Let F be family of functions mapping from Rdd R. q > O. For any > O, with probability least 1 over the random draw of S = (x), ................................ an) D2, the following bound for the test error holds for all F:
2 Asss()+ +3
log(2/8)
Ex (m)
S
1 (.. +
(35)
2n
Let Fr (x (w, 1|l/ S 1) be family of linear functions Rd with bounded weight in no-nomm. The following theorem from Awasthi et al. (2020) bounds the empirical Rademacher complexity of Fpp.
Theorem E.3 (Direct Corollary of Corollary 3 in Awasthi et al. (2020)) For a set of samples (I....., nn the empirical Rademacher complexity of Fi and F2 bounded
1/2
1 (F) < 2 1ogg(2d)
(36)
Y=1
1/2
1 Rs(F2) S |a
(37)
E.2 GENERALIZATION BOUNDS FOR SPARSE LINEAR CLASSIFICATION
Let Dx be the uniform distribution on unif( Let k O(1) be positive odd number. We draw w" as follows: randomly draw the first & coordinates from coordinates to O.
and set all the other
Theorem E.4 With probability at least 1 - over the random draw of the training set S =
(æ1,..., Fn), for any linear classifier Rdd - TR x + (w, x) that maximizes the L--maggin
(ww) = minge[n]
(,,,) on S, the following bound for the test error holds:
2 llg((2)
log(2/8)
EPND (ar 4k
+ 3/
(38)
n
2n
Proof: suffices to consider the case where ||w||| = 1 because rescaling w does not change
the test error. Since arg max (ww), n (w) mww") - By defintion, this implies
35

---

Published as a conference paper at ICLR 2024
ya,( 2 Letting q 2k and applying Theorem E.2, we have
log(2/0)
& 4k (FF 3V
(39)
2n
By Theorem E.3, & $ vn the proof.
2 log(2d)
Combining this with (41) completes
12
Theorem F.5 With probability at least over the random draw of the training set
(Æ1,.... nnd, for any linear classifier f Rd - IR, x + (w,, σ) that maximizes the T22 ma-marnnn
72(w) minie[n]
(ww,w;) on S, the following bound for the test evvor holds:
wwlla l
1io S 4V kd + 3V
log(2/6)
(40)
2n
Proof: suffices consider the case where ||0l||2 because rescaling w does not change the test error. Since arg max 72 (w) (2(22) t By defintion, this implies
ft.() 2 L. Letting q == 1 and applying Theorem E.2, we have
Vk"
2Vk
log(2/8)
E Ex (y)) I 4v (F) + +3
(41)
2n
By Theorem E.3, As(F2) S Vnd Vz Combining this with (41) completes the proof.
36

---

Published as a conference paper at ICLR 2024
F PROOFS FOR MATRIX COMPLETION
In this section, we give the detailed proof of results in Section 3.3.3.
Let Ω = {at = Jik) S k I n and we have n observations Pa. C in 1 S k < n. For convenience, f((0) = (W, P.), and our goal is to minimize the function (5). Since W, is symmetric, it is more convenient replace P. with (P. P, From now on, we let P. = 2 (e,ee, + e)
Pk to denote PæL- Define 0, (U,, Vi). W, U,U, V,V, and
The key step is to show that the loss L defined in (5) satisfies the following two properties: Property F.1 All local minima of C are global.
Property F22 At any saddle point (U,, V,) of C. there is direction (Eu.Ev) such that
vec (E V2
0.
Given these two properties, we can deduce that GF with random initialization can converge to global minimizers almost surely.
Theorem F.3 The function satisfies Properties F.l and F22
Proof:
Recall that
A - vi)²² 201 + vVT).
C(U.V)
Define
(W) Σ
(P W) x²²² 2 W||,,
(42)
1=1
then we have L(U, V) 2 L LUU - -VVT). Since is convex, W is global minimizer of if and only
2 Σ (P.. W v.4 P. - W|L.
0 DVL(WW)
=
(43)
Moreover, for any W Idxd it holds that
min UUT -UT-VVT=W
C(U, V) = L(W).
(44)
now consider some (U,, V.) such that
VCU.,, V.) 0 and VUC V.) 2 0.
To prove the theorem's 's statement, we only need to show that (U., V.) is a global minimizer of L. The subsequent proof is organized into two parts:
Claim 1. W = U.U. - V.V.T T
a global minimizer of é.
Proof of Claim We check that the condition (43) holds at W First the first order condition imply that
VuL(U.,VV)
(P. W - ut) P. P U. =.
(45)
77
Similarly,
Σ (P, -f) 1. V.
(46)
n
37

---

Published as a conference paper at ICLR 2024
Second, for any & R we have
(ViL(U,, v.))
lim (VoL(U,+ +1,, VuC(U,, V.))
--
= =lim 14
(U. + +1) VVV," P.) - -v;) P.(U. +18)
7-0 n i=1
(47)
( W vi) P.U.] +
((P,, Wo) - v.) P.E +2 (U.ET. . P.U.] + XE
where we use (U.ET + +EU, P:)= =2
(U,E, .... since P, is symmetric. Similarly,
(Vic(U., V.)) - " Σ - ((P,, W*) 4i) P.E + 2 2 P.V.V.l + tAE (48)
and
VVVuL(,, V.jE= -- - V,ET, .P.) P.U.
n
(49)
i=1
Let M 1 Ei=1((P
(P., W*) -/) P.. Since V2L(U,, V.) 2 0, for any &u, Ev E Raxd
have
0 S (Ev. (VUL Eu) + (8v. (VEC(U,, V.)) 8v) +2 VVVuCUU,
(Eu,(M + E1) + (Ev, (-) - Σ (U. P.) (v.EV,
8
(V..V,,..) ( (U.ET...))
n
(50)
We now consider two cases:
max (rankU,,. . rankV.) = d. We assume WLOG that rankU, = d, then (45) immediately implies that M + XJ 0. Then -M + AI 2XT and by (46) we have V. case, W is positive semi-definite, and (43) holds since I E 0 | by Lemma F.4.
this
max (rankU,,
< d. In this case, there exists non-zero vectors a, c E Fd such that U.a V.c = 0. For arbitrary vectors b, d TRdd, we set Eu ba" and εv de in (50), so that
|| (M b [|c| (AI M) d20.
Since b, d are arbitrarily chosen, the above implies that
M 0, and A M > 0.
(51)
Since M is symmetric, there exists an orthogonal basis {f: i E [d] of Rd and corre- sponding eigenvalues m,, E [d] such that Mf m: We partition the set [d] into three subsets:
S_ = = [d] -- S+ (i [d] : m, = }}, So = [d] S_ St.
Since (M + U. = 0, we have U.T f: = 0 for all i E S+ U So. As result, we can write
U.U. = Σ aiffff ai E..
Similarly, we can write
V.V. Σ
atffff
ag R.
iJES
38

---

Published as a conference paper at ICLR 2024
Assume WLOG that S_ [t] and S+ {t 1, + 2,....sh), then Hence W Exj=1 aififfz Eij=++1 aij aififf note that the matrices (aij) and (aij)ses are both positive semi-definite, so there exists N 0, [d] and an orthogonal basis (gi i [d] such that
Σ Figfff] f-=1
= Σ bi and 1=1
Σ iJ==t1
azffef]
=
Σ i=++1
bagig,
(52)
and b, 0 and g: = fi for i s. now let
G = (g= : [d)), G2 (a1)))l00gg:iiee
and B diag|b:| : E [d]
then both 11 and G2 are orthonormal matrices and we have W = BG This gives a singular value decomposition of W By Lemma E.4 write down the explicit form of 0|| W* as follows:
1|
(120))
9:g: E: | S 1 and Eg. = 0, 0.
iET
Here T {i E [d] b, + 0} C [s]. We choose
E 1 M + Σ - 1) (120) 9:9:
SE
-
f.f
f.f.
Emfift
(1) 9:9.
xm++1
-X
9:9
+
Σ
9:9.
Σ
mggig?
(-1) g:g.
1=1
|=H+
1>8
iET
Σ 1E(11-T
(-1)
(i>t)
9:9;
-1 1 Σ
m4g:9;
where the third equation holds because of the definition (52) and g: = f: when i > 8. Since by (51) we have m, E I-A, A] for all i E [d]. the above expression of E immediately implies that || [ < 1. Moreover, we obviously have Eg: = when E TT (since {g:) is orthogonal). Hence, have 0 M A0 W*
To summarize, we have shown that (43) always holds at W*, Therefore, W is a global minimizer of È which concludes the proof of Claim
Claim 2. (U., V.) solves the problem
minimize
+ s.t. UUT VV = W
(53)
Proof of Claim We will use the notations in the proof of Claim for convenience. note that
-U.12=-trUa U. b
and similarly
V. 12 b.
1=
Thus
IU_ 13. + || 112. = Σi b, = W
W| ||. On the other hand, for any (U, V) satisfying UUT
VVT = W*, we have
[W'l - VV |1. < LU 11. + +||V U
|
39

---

Published as a conference paper at ICLR 2024
Hence (U,, V.) is minimizer of (53), as desired.
We can prove that (Uz, V.) is global minimizer of L, which completes the proof of Theo- rem F.3. Indeed, for any (U, V) we have
L(U, V) 2 LLUUT - - 2 L(W")= L(U., V.).
The last step follows from Claim 2 and (44).
Lemma F.4 (Watson, 1992, Example 2) Let = LER its singular value decomposition, then we have
0| A|l {LR* + E E < 1, L E and ER =0].
It is well-knows that under cerntain regularity conditions on the ground-truth matrix X", the global minimizer of the convex problem (42) is close the ground-truth X Here we present a version of this result adapted from (Candes and Plan, 2010, Theorem 7).
Theorem F,5 Suppose that rank x*) = = 0(1) X = -V-- -Vx-Ex-Vy is (a version of) its SVD, and each row of Vx has Lox norm bounded by / then if the number of observed entries satisfies 2 pdd log" d, then we have that | X* || vAlxX*Il log d with probability d-3
40


Reviews:
[{"review_id": "BjFXMPDLvE", "reviewer": "Reviewer_XQUe", "summary": "This paper studies grokking theoretically in homogeneous neural networks with large initialization scale and small weight decay. It shows that in both classification and regression settings, there are early and late implicit biases. Particularly, it shows that in the early phase the model corresponds to a kernel model and later it can escape it. It suggests that grokking happens when the solution given by the early implicit bias generalizes poorly while the solution given by the late bias generalizes well. It is also shown that the converse phenomenon can happen which they call misgrokking. \n\nAs an example, the paper focuses on a simple diagonal linear neural net for classification and shows that the early and late bias resp. correspond to minimization of L2 and L1 margin. They also identify these biases for a matrix completion model. \n\nThe paper further provides experiments showing that grokking can happen even in the absence of weight decay.", "questions": "- Q1. What are the limitations of the analysis done in this paper? Ignoring the proof difficulties, do you think the change of implicit biases can be observed and explain grokking in more common settings?\n- Q2. What would be the effects of stochasticity in GD (SGD) and large step size (distancing from gradient flow)?\n- Q3. Can you further explain Remark 3.11?\n- Q4. Can you further explain the learning rate and time scaling for Figure 2? (E.g., the $\\log(t)$ has often very large values.)\n- Q5. Is there any proof provided for Theorem 4.1? \n- Q6. Can you further clarify why proving $ \\|\\frac 1\\alpha\n \\exp(\\lambda t) \\theta(t) - \\overline{\\theta}_{\\mathrm{init}} \\| \\leq \\epsilon$   for all $t \\leq \\min \\{\\frac{1}{\\lambda} \\log \\alpha + \\cdots + \\Delta T, T_m \\} $ would imply $T_m \\geq \\frac 1 \\lambda \\log \\alpha + \\cdots + \\Delta T$ in proofs of Lemma A.5 and B.2?\n- Q7. What is $q(\\theta)$ in Appendix A.2? Also $r_i$'s?\n\n## Minor Questions/Remarks\n- Q8. In Figure 1, what is the norm presented? (Whether it's L1/L2 and if it's for $w$ or $u,v$?)\n- Q9. I think the analyses require the data to be linearly separable for the classification setting? (Also, the model should be able to interpolate in the regression setting.) Although these assumptions are reasonable for overparametrized model, it would be beneficial if they were stated more clearly in the main text. \n- Q10. In the paragraph before Corollary 3.5, isn't the kernel feature $(2x, -2x)$?\n- R1. In the paragraph of Empirical Validation: Grokking, it would be nice if some references for $O(d)$ and $O(k\\log d)$ sample complexities were provided. Further, if some references/theoretical results for poor generalization of L2 margin are provided, it would complete the theoretical part of grokking for classification.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is quite rigorous and provides a provable instance of grokking for simple models. Furthermore, it attributes grokking to a transition from the kernel to the rich regime. \n- The paper does the analyses for both classification and regression settings, further it shows the possibility of misgrokking. \n- Empirically, it's shown that grokking can happen without weight decay but more slowly.", "weaknesses": "- There is very limited discussion on the possible extensions and limitations of this model. Moreover, results are in the case that the initialization scale goes to infinity which doesn't show (at least immediately) if these observations can be observed in the more common settings. \n- Additional experiments can clarify other potential limitations of the analysis (e.g., gradient flow). \n\nSee the questions for more details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3eifkwiVSA", "reviewer": "Reviewer_oyWf", "summary": "The paper studies an empirical phenomenon called \"Grokking\" where-in neural networks can display delayed generalization under certain conditions. The paper puts forth the hypothesis that the observed memorization and eventual transition to generalization can be attributed to the optimization occurring two different inductive biases that are separated by a transition at a certain time (step). The paper studies training of neural network under large initialization and weight decay and suggests that for the first block of time the training occurs in a \"kernel regime\" before transitioning to \"rich regime\". The paper studies behavior under classification and regression settings using diagonal linear networks as a concrete example", "questions": "- In Section 3.1.3 on page 5 of the draft, the paper claims that implicit biases imply transition in test accuracy. As noted in the weakness, this is not supported by any analysis (bounds etc). \n- No empirical results are provided that helps the reader crystallize their understanding of the work in terms of when the transition might occur in practice. Can the authors add more examples in text with datasets commonly used in ``Grokking'' literature?\n- One point that the paper makes is about identifying a transition time from memorization to generalization and attributes it to different inductive biases. This is not true when weight decay = 0 which is a case handled that the paper talks about as well. This suggestion weakens about the importance of the first phase as it may not even be necessary to understand ``Grokking''", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper attempts to explain \"Grokking\" using analytical tools. As an empirical researcher, I found this to be a good reminder that theory can shed light on even perplexing phenomena ``Grokking''\n- The paper implicitly summarizes the various conditions under which ``Grokking'' has been observed and reported in practice\n- The paper's hypothesis of two separate indicative biases separated by transition at a certain time makes intuitive sense\n- The paper empirically constructs examples that empirically  demonstrate both Grokking and mis-Grokking by manipulating the inductive biases while creating datasets", "weaknesses": "The biggest concern that I have is that I do not quite see the core claim os \"sharp transition in test accuracy\" being supported by any of the analysis made in the paper. The paper states theorems in the main paper that show that two different inductive biases are at work and provide corollaries for diagonal linear networks. However, I do not see any statements, analysis or proofs in the main paper for generalization risks (or make generalization guarantees).\n\nWithout the above, the best one can conclude is that the models can fit the training data and reach perfect training accuracy eventually but this line of work has already been done by Moroshko et. al (Moroshko). I may have missed something obvious so I look forward to the rebuttal to improve my understanding of the paper\n\nHowever, at this point, I am afraid I can't really support accepting this paper.\n\n[Moroshko] https://arxiv.org/abs/2007.06738", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nyuvX24gcv", "reviewer": "Reviewer_vAjC", "summary": "The authors seek to explain the phenomenon of grokking highlighted by Power, Burda, Edwards, Babuschkin, and Misra in arXiv 2201.02177, focusing on a theoretical analysis, and paying special attention to sharpness of the transition from memorisation to generalisation.  The central thesis in this paper, which is summarised in its title, is that grokking occurs due to different implicit biases in early and late training.  Several theoretical results in support of this are proved for training by gradient flow homogeneous neural networks with large initialisations and small weight decay, for classification as well as regression tasks.  The pattern is that general theorems are proved first, and then corollaries are derived that explain grokking in specific settings, such as classification using diagonal linear networks, or low-rank symmetric matrix completion.  In the final section, the authors consider other scenarios where grokking occurs; they provide numerical experiments without weight decay for learning modular addition using two-layer ReLU networks, and they show how theoretical results from Li, Wang, and Arora in ICLR 2021 confirm grokking when training two-layer diagonal linear networks by label noise SGD with a sufficiently small learning rate.", "questions": "In the first sentence in \"Our Contributions\", it is claimed that homogeneous networks allow the ReLU activation, however homogeneity is then assumed together with $\\mathcal{C}^2$-smoothness (in Assumption 3.1), which I believe disallows the ReLU activation due to its non-differentiability at 0?\n\nCode was not submitted with the paper, only the appendix with proofs.  Are you able to supplement the paper with code, so that readers are able to reproduce your numerical experiments?\n\nWhat is the purpose of the section on sharpness reduction with label noise (Section 4.2), since it mostly consists of statements of results from the paper of Li, Wang, and Arora in ICLR 2021, i.e. it does not seem to contain much new material?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper seems to be a first one to succeed in proving rigorously that grokking occurs in relatively general settings, and in showing interesting upper bounds on the length of the transition from memorisation to generalisation.\n\nThe explanation of grokking in terms of different implicit biases in early and late phases of training is elegant, and the restrictions to homogeneous neural networks, large initialisation, small weight decay, and gradient flow are reasonable and motivated by the previous empirical work of Liu, Michaud, and Tegmark in ICLR 2023.  Moreover, the versatility of the explanation in terms of the implicit biases is demonstrated by the identification of \"misgrokking\" in this paper, a dual phenomenon in which the test accuracy transitions from very high to very low late in the training.\n\nGrokking and misgrokking as indicated by the theory are verified by numerical experiments for classification using two-layer diagonal linear networks.\n\nFull proofs of the theoretical results are provided in the appendix.  The biggest and most novel are those of Theorems 3.3 and 3.7, which show that the small weight decay does not spoil the kernel regime and its implicit bias early in the training for long enough so that a relatively short transitional phase follows.", "weaknesses": "The paper comes across as dense and can be difficult to read.  One reason is that a number of settings for theory and experiments are considered.  Another is that the statements of many of the theoretical results are relatively complex, with several quantifications and subtle details, and often without much explanation.  A minor comment is that I think that writing \"for all\" or \"for every\" is clearer than \"for any\".\n\nThe theoretical results about grokking for regression tasks, and in particular for overparameterised matrix completion, are not illustrated experimentally.\n\nThe conclusion of the paper (Section 5) does not seem to indicate any directions for future work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)