Original: O MNIGROK : G ROKKING B EYOND A LGORITHMIC D ATA
Revised: O MNIGROK : G ROKKING B EYOND A LGORITHMIC D ATA

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Grokking, the unusual phenomenon for algorithmic datasets where generalization happens long after overfitting the training data, has remained elusive.
Revised: Grokking, the unusual phenomenon for algorithmic datasets where generalization happens long after overfitting the training data, has remained elusive.

Original: We aim to understand grokking by analyzing the loss landscapes of neural networks, identifying the mismatch between training and test losses as the  cause for grokking.
Revised: We aim to understand grokking by analyzing the loss landscapes of neural networks, identifying the dependence of the generalization gap on model weight norm as a cause of grokking.

Original: We refer to this as the "LU mechanism" because training and test losses (against model weight norm) typically resemble "L" and "U", respectively.
Revised: We refer to this as the "LU mechanism" because training and test losses (against model weight norm) typically resemble "L" and "U", respectively.

Original: Thissimple mechanism cannicely explain many aspects of grokking: data size dependence, weight decay dependence, the emergence of representations, etc.
Revised: This mechanism can explain many aspects of grokking: data size dependence, weight decay dependence, the emergence of representations, etc.

Original: Guided by the intuitive picture, we are able to induce grokking on tasks involving images, language and molecules, although the grokking signals are sometimes less dramatic.
Revised: Guided by the intuitive picture, we are able to induce grokking on tasks involving images, language and molecules, although the grokking signals are sometimes less dramatic.

Original: We attribute the dramatic nature of grokking for algorithmic datasets to representation learning.
Revised: We attribute the dramatic nature of grokking for algorithmic datasets to representation learning.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Generalization lies at the heart of machine learning.
Revised: Generalization lies at the heart of machine learning.

Original: A good machine learning model should arguably be able to generalize fast, and behave in a smooth/predictable way under changes of (hyper)parameters.
Revised: A good machine learning model should arguably be able to generalize fast, and behave in a smooth/predictable way under changes of (hyper)parameters.

Original: Grokking, the phenomenon where the model generalizes long after overfitting the training set, has raised interesting questions after it was observed on algorithmic datasets by Power et al.
Revised: Grokking, the phenomenon where the model generalizes long after overfitting the training set, has raised interesting questions after it was observed on algorithmic datasets by Power et al.

Original: Q

Revised: Q

Original: The origin of grokking : Why is generalization much delayed after overfitting?
Revised: The origin of grokking : Why is generalization much delayed after overfitting?

Original: Q

Revised: Q

Original: The prevalence of grokking : Can grokking occur on datasets other than algorithmic datasets?
Revised: The prevalence of grokking : Can grokking occur on datasets other than algorithmic datasets?

Original: This paper aims to answer these questions by analyzing neural loss landscapes:
Revised: This paper aims to answer these questions by analyzing neural loss landscapes:

Original: A1 Grokking can result from a mismatch between training and test loss against model weight norm.
Revised: A1 Grokking can result from a mismatch between training and test loss against model weight norm.

Original: Specifically, (reduced) training and test losses plotted against model weight norm resemble "L" and "U", respectively, as shown in Figure 1b.
Revised: Specifically, (reduced) training and test losses plotted against model weight norm resemble "L" and "U", respectively, as shown in Figure 1b.

Original: We refer to this phenomenon as the "LU mechanism", which we elaborate on in Section 2 and 3.
Revised: We refer to this phenomenon as the "LU mechanism", which we elaborate on in Section 2 and 3.

Original: A

Revised: A

Original: Yes .
Revised: Yes .

Original: Indeed, we demonstrate grokking for a wide range of machine learning tasks in Section 4, including image classification sentiment analysis and molecule property prediction
Revised: Indeed, we demonstrate grokking for a wide range of machine learning tasks in Section 4, including image classification, sentiment analysis and molecule property prediction.

Original: Grokking signals observed for these tasks are usually less dramatic than for algorithmic datasets, which we attribute to representation learning in Section 5.
Revised: Grokking signals observed for these tasks are usually less dramatic than for algorithmic datasets, which we attribute to representation learning in Section 5.

Original: Partial answers to Q1 are provided in recent studies: Liu et al.
Revised: Partial answers to Q1 are provided in recent studies: Liu et al.

Original: (2022) attribute grokking to the slow formation of good representations, Thilak et al.
Revised: (2022) attribute grokking to the slow formation of good representations, Thilak et al.

Original: (2022) attempts to link grokking to the slingshot mechanism of adaptive optimizers, and Barak et al.
Revised: (2022) attempts to link grokking to the slingshot mechanism of adaptive optimizers, and Barak et al.

Original: (2022) uses Fourier gap to describe hidden progress.
Revised: (2022) uses Fourier gap to describe hidden progress.

Original: This paper aims to understand grokking through the lens of neural loss landscapes.
Revised: This paper aims to understand grokking through the lens of neural loss landscapes.

Original: Our landscape analysis is able to explain many aspects of grokking: data size dependence, weight decay dependence, emergence of representations, etc.
Revised: Our landscape analysis is able to explain many aspects of grokking: data size dependence, weight decay dependence, emergence of representations, etc.

Original: The paper is organized as follows:
Revised: The paper is organized as follows:

Original: In Section 2, we review background on generalization, and introduce the LU mechanism .
Revised: In Section 2, we review background on generalization, and introduce the LU mechanism .

Original: In Section 3, we show how the LU mechanism leads to grokking for a toy teacher-student setup.
Revised: In Section 3, we show how the LU mechanism leads to grokking for a toy teacher-student setup.

Original: In Section 4, we show that the intuition gained from the toy problem can transfer to realistic datasets (MNIST, IMDb reviews and QM9), for which we also observe grokking, although in a slightly non-standard setup where it is relatively weak.
Revised: In Section 4, we show that the intuition gained from the toy problem can transfer to realistic datasets (MNIST, IMDb reviews and QM9), for which we also observe grokking, although in a slightly non-standard setup where it is relatively weak.

Original: In Section 5, we discuss why grokking is more dramatic for algorithmic datasets than on others (e.g., MNIST), by comparing their loss landscapes.
Revised: In Section 5, we discuss why grokking is more dramatic for algorithmic datasets than on others (e.g., MNIST), by comparing their loss landscapes.

Original: We review related work in Section 6 and summarize our conclusions in Section 7.
Revised: We review related work in Section 6 and summarize our conclusions in Section 7.

Original: w > w c w w c
Revised: Code is available at https://github.com/KindXiaoming/Omnigrok .

Original: w  < w c large init small init fast fast s l o w ( g r o k k i n g ) generalizable solutions overfitting solutions weight space ( [Figure] w c weight norm w
Revised: w > w c w w c w < w c large init small init fast fast s l o w ( g r o k k i n g ) generalizable solutions overfitting solutions weight space (a) w c weight norm w

Original: L o ss f a s t s l o w ( g r o k k i n g ) t
Revised: L o ss f a s t s l o w ( g r o k k i n g ) t

Original: e s t l o ss t r a i n i n g l o s s l a r g e r e g s m a ll r
Revised: e s t l o ss t r a i n i n g l o s s l a r g e r e g s m a ll r

Original: e g no reg large init small init
Revised: e g no reg large init small init

Original: f a s t
Revised: f a s t

Original: G o l d il o c k s z
Revised: G o l d il o c k s z

Original: o n e ( g e n e r a li z a t i o n ) (
Revised: o n e ( g e n e r a li z a t i o n ) (b)

Original: [Figure]

Original:  T HE LU MECHANISM FOR GROKKING
Revised:  T HE LU MECHANISM FOR GROKKING

Original: Weight norm and reduced loss Letting w denote the weights of a model, any function f ( w ) (e.g, train/test loss/accuracy) depends on both the weight norm
Revised: Weight norm and reduced loss Letting w denote the weights of a model, any function f ( w ) (e.g, train/test loss/accuracy) depends on both the weight norm

Original: w ≡ || w || 2 and the angular direction ˆ w ≡ w /w .
Revised: w ≡ || w || 2 and the angular direction ˆ w ≡ w /w .

Original: Similar to Fort and Scherlis (2019), we define a reduced function ˜ f
Revised: Similar to Fort and Scherlis (2019), we define a reduced function ˜ f

Original: ( w ) by minimizing training loss l train ( w ) over angular directions, i.e.,˜ f
Revised: ( w ) by minimizing training loss l train ( w ) over angular directions, i.e.,

Original: [Equation]
Revised: [Equation]

Original: f
Revised: f

Original: [Equation]
Revised: [Equation]

Original: w ∗ ( w )) , where w ∗ ( w ) ≡ argmin 
Revised: w ∗ [Equation] , where w ∗ ( w ) ≡ argmin [Equation]

Original: [Equation]
Revised: [Equation]

Original: In this paper, we set f as train/test loss/error, but it also applies to other metrics of interest.
Revised: In this paper, we set f as train/test loss/error, but it also applies to other metrics of interest.

Original: In practice, we perform the constrained minimization by rescaling the model weights back to their original norm after each unconstrained optimization step.
Revised: In practice, we perform the constrained minimization by rescaling the model weights back to their original norm after each unconstrained optimization step.

Original: We will see that this reduced 1D loss landscape, which is easy to visualize, captures important features related to grokking.
Revised: We will see that this reduced 1D loss landscape, which is easy to visualize, captures important features related to grokking.

Original: Throughout the paper, our model is initializaed by multiplying a factor α ≡ w/w 0 to the standard initialization 1 , where w 0 and w are the weight norm of the network before and after multiplying  α .
Revised: Throughout the paper, our model is initialized by multiplying a factor α ≡ w/w 0 to the standard initialization 1 , where w 0 and w are the weight norm of the network before and after multiplying by α , respectively.

Original: LU mechanism Although the loss landscapes of neural networks are nonlinear, Fort and Scherlis (2019) reveal a simple landscape picture: There is a spherical shell in the weight space (the "Goldilocks" zone), where generalization is better than outside this zone.
Revised: LU mechanism Although the loss landscapes of neural networks are nonlinear, Fort and Scherlis (2019) reveal a simple landscape picture: There is a spherical shell in the weight space (the "Goldilocks" zone), where generalization is better than outside this zone.

Original: We illustrate the Goldilocks zone as the green area with average radius w c in Figure 1a; the green stars are the generalizing solutions.
Revised: We illustrate the Goldilocks zone as the green area with average radius w c in Figure 1a; the green stars are the generalizing solutions.

Original: The test loss is thus higher either both when w > w c and w < w c , forming a U-shape against w in Figure 1b (gray curve).
Revised: The test loss is thus higher either both when w > w c and w < w c , forming a U-shape against w in Figure 1b (gray curve).

Original: By contrast, the training loss has an L-shape against weight norm .
Revised: By contrast, the training loss has an L-shape against weight norm .

Original: There are many solutions which overfit training data for w > w c , but high training losses are incurred for w < w c .
Revised: There are many solutions which overfit training data for w > w c , but high training losses are incurred for w < w c .

Original: This corresponds to the L-shaped curve seen in Figure 1b (orange curve, no regularization).
Revised: This corresponds to the L-shaped curve seen in Figure 1b (orange curve, no regularization).

Original: In summary, the (reduced) training loss and test loss are L-shaped and U-shaped against weight norm, respectively, which we will refer to as the LU mechanism throughout the paper.
Revised: In summary, the (reduced) training loss and test loss are L-shaped and U-shaped against weight norm, respectively, which we will refer to as the LU mechanism throughout the paper.

Original: It is well known in statistics that generalization error has a "U" shape against model capacity, which is usually attributed to the bias-variance trade-off.
Revised: It is well known in statistics that generalization error has a "U" shape against model capacity, which is usually attributed to the bias-variance trade-off.

Original: Although this common wisdom was challenged by the observation of double descent (Nakkiran et al., 2021), the "U" curve can be recovered from a double descent simply by changing the x-axis from the number of model parameters N to the 2-norm of model parameters w ≡ || w || 2 , at least for linear regression (Ng and Ma, 2022).
Revised: Although this common wisdom was challenged by the observation of double descent (Nakkiran et al., 2021), the "U" curve can be recovered from a double descent simply by changing the x-axis from the number of model parameters N to the 2-norm of model parameters w ≡ || w || 2 , at least for linear regression (Ng and Ma, 2022).

Original: Although the LU mechanism may remind readers of related phenomena (Schoenholz et al., 2016; Yang and Schoenholz, 2017; Nakkiran et al., 2021), their setups are not exactly the same as ours.
Revised: Although the LU mechanism may remind readers of related phenomena (Schoenholz et al., 2016; Yang and Schoenholz, 2017; Nakkiran et al., 2021), their setups are not exactly the same as ours.

Original: More importantly, our focus and contribution is to understand grokking, a brand new generalization puzzle.
Revised: More importantly, our focus and contribution is to understand grokking

Original: Grokking dynamics We identify the "LU mechanism" as the cause of grokking.
Revised: Grokking dynamics We identify the "LU mechanism" as the cause of grokking.

Original: If the weight norm is initialized to be large (e.g., the black square in the w > w c region), the model first quickly moves to a nearby overfitting solution by minimizing the training loss.
Revised: If the weight norm is initialized to be large (e.g., the black square in the w > w c region), the model first quickly moves to a nearby overfitting solution by minimizing the training loss.

Original: Without any regularization, the model will stay where it is, because the gradient of the training loss is almost zero along the valley of overfitting solutions, so generalization does not happen.
Revised: Without any regularization, the model will stay where it is, because the gradient of the training loss is almost zero along the valley of overfitting solutions, so generalization does not happen.

Original: Fortunately, there are usually explicit and/or implicit regularizations that can drive the weight vector towards the Goldilocks zone w ≈ w c .
Revised: Fortunately, there are usually explicit and/or implicit regularizations that can drive the weight vector towards the Goldilocks zone w ≈ w c .

Original: When the regularization magnitude is non-zero but small, the radial motion can be (arbitrarily) slow.
Revised: When the regularization magnitude is non-zero but small, the radial motion can be (arbitrarily) slow.

Original: If weight decay is the only source of regularization, and training loss is negligible after overfitting, then weight decay γ causes w ( t ) ≈ exp( − γt ) w 0 , when w 0 > w c , so it takes time t ≈ ln( w 0 /w c ) /γ
Revised: If weight decay is the only source of regularization, and training loss is negligible after overfitting, then weight decay γ causes w ( t ) ≈ exp( − γt ) w 0 , when w 0 > w c , so it takes time t ≈ ln( w 0 /w c ) /γ

Original: ∝ γ − 1 to generalize.
Revised: ∝ γ − 1 to generalize.

Original: A small γ results in a huge generalization delay (i.e., grokking).
Revised: A small γ results in a huge generalization delay (i.e., grokking).

Original: The dependence on regularization magnitudes is illustrated in Figure 1b: no generalization at all happens for γ = 0 , small γ leads to slow generalization (grokking), and large γ leads to faster generalization 2 .
Revised: The dependence on regularization magnitudes is illustrated in Figure 1b: no generalization at all happens for γ = 0 , small γ leads to slow generalization (grokking), and large γ leads to faster generalization 2 .

Original: The above analysis only applies to large initializations w > w c .
Revised: The above analysis only applies to large initializations w > w c .

Original: Small initializations w < w c can always generalize fast 3 , regardless of regularization.
Revised: Small initializations w < w c can always generalize fast 3 , regardless of regularization.

Original: Why isn’t grokking commonly observed?
Revised: Why isn’t grokking commonly observed?

Original: The standard initialization schemes typically initialize w no larger than w c .
Revised: The standard initialization schemes typically initialize w no larger than w c .

Original: However, if we increase initialization scales (explicitly or implicitly), grokking can appear.
Revised: However, if we increase initialization scales (explicitly or implicitly), grokking can appear.

Original: In Section 3 and 4, we find that explicitly increasing initialization weight norm can induce grokking.
Revised: In Section 3 and 4, we find that explicitly increasing initialization weight norm can induce grokking.

Original: In Section 5, we argue for algorithmic datasets because (shown in Figure 6d) [Equation] i.e., a proper initialization for a bad representation is effectively too large for a good representation, leading to grokking.
Revised: In Section 5, we argue for algorithmic datasets because (shown in Figure 6d) w c (bad representation) > w c (good representation) , (2) i.e., a proper initialization for a bad representation is effectively too large for a good representation, leading to grokking.

Original: Take the addition (base p ) for example: with the good (linear) representation or a bad (random) representation, the decoder needs to learn to classify O ( p ) or O ( p 2 ) examples, respectively [Figure]
Revised: Take the addition (base p ) for example: with the good (linear) representation or a bad (random) representation, the decoder needs to learn to classify O ( p ) or O ( p 2 ) examples, respectively.

Original:  G ROKKING FOR A TEACHER - STUDENT SETUP
Revised:  G ROKKING FOR A TEACHER - STUDENT SETUP

Original: [Figure]
Revised: [Figure]

Original: To illustrate how the LU mechanism results in grokking, we employ a toy teacher-student setup.
Revised: To illustrate how the LU mechanism results in grokking [Figure] we employ a toy teacher-student setup.

Original: The teacher and the student share the same architecture (a 5-100-100-5 MLP with tanh activation), but are initialized with different seeds.
Revised: The teacher and the student share the same architecture (a 5-100-100-5 MLP with tanh activation), but are initialized with different seeds.

Original: The student network is initialized with the standard initialization (the default one in PyTorch) but each weight is rescaled by the same factor α ≡ w/w 0 , where wand w are the weight norm of the student network before and after rescaling.
Revised: The student network is initialized with the standard initialization (the default one in PyTorch) but each weight is rescaled by the same factor α ≡ w/w 0 , where wand w are the weight norm of the student network before and after rescaling.

Original: The teacher network is initialized standardly, i.e., α teacher = 1 .
Revised: The teacher network is initialized standardly, i.e., α teacher = 1 .

Original: Inputs and outputs have dimensions d in = 5 and d out = 5 , respectively.
Revised: Inputs and outputs have dimensions d in = 5 and d out = 5 , respectively.

Original: We generate N train = 100 training and N test = 100 test samples by first drawing inputs from the standard Gaussian distribution N (0 , I d in × d in ) , and then feed the input data to the teacher to generate output labels.
Revised: We generate N train = 100 training and N test = 100 test samples by first drawing inputs from the standard Gaussian distribution N (0 , I d in × d in ) , and then feed the input data to the teacher to generate output labels.

Original: The student network is trained with the Adam optimizer (learning rate 3 × 10 − 4 ) for 10 5 steps.
Revised: The student network is trained with the Adam optimizer (learning rate 3 × 10 − 4 ) for 10 5 steps.

Original: LU landscapes Firstly, we compute the reduced losses by minimizing the training loss (excluding weight decay) while constraining the weight norm of the student network to be constant.
Revised: LU landscapes Firstly, we compute the reduced losses by minimizing the training loss (excluding weight decay) while constraining the weight norm of the student network to be constant.

Original: We assume the converging point after training as the global minimum on the spherical surface 4 , which explicitly defines the reduced losses ˜ l train ( α ) and ˜ l test ( α ) .
Revised: We assume the converging point after training as the global minimum on the spherical surface 4 , which explicitly defines the reduced losses ˜ l train ( α ) and ˜ l test ( α ) .

Original: As shown in Figure 2a, ˜ l test ( α ) first decreases and then increases as α increases, displaying a U-shape with a minimum at α ≈ 1 .
Revised: As shown in Figure 2a, ˜ l test ( α ) first decreases and then increases as α increases, displaying a U-shape with a minimum at α ≈ 1 .

Original: By contrast, ˜ l train ( α ) decreases when α < 1 and remains flat near zero when α ≥ 1 , forming an L-shape.
Revised: By contrast, ˜ l train ( α ) decreases when α < 1 and remains flat near zero when α ≥ 1 , forming an L-shape.

Original: When weight decay γ is present, the training landscape becomes ˜ l train ( α, γ ) =
Revised: When weight decay γ is present, the training landscape becomes ˜ l train ( α, γ ) =

Original: ˜ l train ( α ) + γα 2 C 2 where C is the average parameter magnitude determined by the standard initialization.
Revised: ˜ l train ( α ) + γα 2 C 2 where C is the average parameter magnitude determined by the standard initialization.

Original: Training dynamics
Revised: Training dynamics

Original: Our problem is a regression task, but we can imitate the behavior of a classification task by manually setting a threshold β = 0 .
Revised: Our problem is a regression task, but we can imitate the behavior of a classification task by manually setting a threshold β = 0 .

Original: 01 and defining a sample to be correctly “classified" if the prediction error is less than β .
Revised: 01 and defining a sample to be correctly “classified" if the prediction error is less than β .

Original: We study the dynamics of training and test accuracy.
Revised: We study the dynamics of training and test accuracy.

Original: Note that this is the normal training setup where the weight norm is not constrained, although with two non-standard initializations α = 0 .
Revised: Note that this is the normal training setup where the weight norm is not constrained, although with two non-standard initializations α = 0 .

Original: 5 (small) and α = 2 . 

Revised: 5 (small) and α = 2 .

Original: (large), and three weight decays γ = 

Revised: (large), and three weight decays γ =

Original: (no reg), γ = 0 .
Revised: (no reg), γ = 0 .

Original: 03 (small reg) and γ = 1 (large reg).
Revised: 03 (small reg) and γ = 1 (large reg).

Original: As shown in Figure 2b (bottom), small initialization runs always generalize fast regardless of regularization.
Revised: As shown in Figure 2b (bottom), small initialization runs always generalize fast regardless of regularization.

Original: Large initialzation runs (top) dependend on weight decay: no regularization fails to generalize, small regularization generalizes slowly (grokking), while large regularization generalizes faster.
Revised: Large initialzation runs (top) dependend on weight decay: no regularization fails to generalize, small regularization generalizes slowly (grokking), while large regularization generalizes faster.

Original: For the large initialization α = 2 . 0 , we do a finer sweep of γ in [0 . 03 , 1] .
Revised: For the large initialization α = 2 . 0 , we do a finer sweep of γ in [0 . 03 , 1] .

Original: We compute the number of steps and weight norm w when training or test accuracy reaches 95%.
Revised: We compute the number of steps and weight norm w when training or test accuracy reaches 95%.

Original: As shown in Figure 2c, the time (number of steps) to reach 95% training accuracy is independent of weight decay γ , while the time to reach 95% test accuracy is inversely proportional to the weight decay, as we derived above for the LU mechanism.
Revised: As shown in Figure 2c, the time (number of steps) to reach 95% training accuracy is independent of weight decay γ , while the time to reach 95% test accuracy is inversely proportional to the weight decay, as we derived above for the LU mechanism.

Original:  O MNIGROK : G ROKKING FOR MORE INTERESTING TASKS
Revised:  O MNIGROK : G ROKKING FOR MORE INTERESTING TASKS

Original: We now analyze loss landscapes and search for grokking for several more interesting datasets, and see that the insights obtained from our toy model can transfer to these datasets.
Revised: We now analyze loss landscapes and search for grokking for several more interesting datasets, and see that the insights obtained from our toy model can transfer to these datasets.

Original: We report the main results here, with experiment details included in Appendix A.
Revised: We report the main results here, with experiment details included in Appendix A.

Original: Image classification We visualize loss landscapes of MNIST (Deng, 2012) to verify the LU mechanism, and study the dependence on training data size.
Revised: Image classification We visualize loss landscapes of MNIST (Deng, 2012) to verify the LU mechanism, and study the dependence on training data size.

Original: Similar to the teacher-student case, we reduce losses and errors (one minus accuracy) to two variables (weight norm w and data size N ) by minimizing over angular directions of weights, i.e.,
Revised: Similar to the teacher-student case, we reduce losses and errors (one minus accuracy) to two variables (weight norm w and data size N ) by minimizing over angular directions of weights, i.e.,

Original: [Equation]
Revised: [Equation]

Original: shown in Figure 3 (a)(b).
Revised: shown in Figure 3 (a)(b).

Original: The reduced loss landscape reveals three things: (1) Larger initializations lead to grokking.
Revised: The reduced loss landscape reveals three things: (1) Larger initializations lead to grokking.

Original: Point A in Figure 3 corresponds to the standard initialization ( α = 1) , which has low training and test errors, hence no grokking.
Revised: Point A in Figure 3 corresponds to the standard initialization ( α = 1) , which has low training and test errors, hence no grokking.

Original: When increasing the weight norm from A to
Revised: When increasing the weight norm from A to

Original: B , training error is seen to remain low while test error rises.
Revised: B , training error is seen to remain low while test error rises.

Original: To generalize, weight decay must be in place to bring the weight norm down, leading to grokking if weight decay is small.
Revised: To generalize, weight decay must be in place to bring the weight norm down, leading to grokking if weight decay is small.

Original: (2) Larger datasets lead to de-grokking.
Revised: (2) Larger datasets lead to de-grokking.

Original: Comparing B and C in Figure 3, C is seen to have larger training size than B and lower test error.
Revised: Comparing B and C in Figure 3, C is seen to have larger training size than B and lower test error.

Original: Larger data size N makes the Goldilocks zone broader, reducing or eliminating grokking even for large weight initializations.
Revised: Larger data size N makes the Goldilocks zone broader, reducing or eliminating grokking even for large weight initializations.

Original: (3) Critical data size can be defined.
Revised: (3) Critical data size can be defined.

Original: As reported in Power et al.
Revised: As reported in Power et al.

Original: Liu et al.
Revised: Liu et al.

Original: (2022), we see that there exists a critical training set size below which generalization is impossible.
Revised: (2022), we see that there exists a critical training set size below which generalization is impossible.

Original: The effective theory analysis in Liu et al.
Revised: The effective theory analysis in Liu et al.

Original: (2022) only applies to algorithmic datasets, but not to other datasets with unknown optimal representations.
Revised: (2022) only applies to algorithmic datasets, but not to other datasets with unknown optimal representations.

Original: The loss landscape analysis presented is this work can apply to all supervised-learning tasks.
Revised: The loss landscape analysis presented is this work can apply to all supervised-learning tasks.

Original: As shown in
Revised: As shown in

Original: Figure 3 (b), the contours of constant test error are thumb-like, and the tip of the thumb determines the minimum amount of data required for generalization.
Revised: Figure 3 (b), the contours of constant test error are thumb-like, and the tip of the thumb determines the minimum amount of data required for generalization.

Original: Guided by the landscape analysis, we make two nonstandard decisions to induce grokking on MNIST: (1) we reduce the size of the training set from 60k to 1k samples (by taking a random subset) and (2) we increase the scale of the weight initialization distribution (by multiplying the initial weights, sampled with Kaiming uniform initialization, by a constant α > 1 ).
Revised: Guided by the landscape analysis, we make two nonstandard decisions to induce grokking on MNIST: (1) we reduce the size of the training set from 60k to 1k samples (by taking a random subset) and (2) we increase the scale of the weight initialization distribution (by multiplying the initial weights, sampled with Kaiming uniform initialization, by a constant α > 1 ).

Original: With these modifications to the
Revised: With these modifications to the

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: training set size and initialization scale, we train a depth-3 width-200 MLP with ReLU activations with the AdamW optimizer using MSE loss with one-hot targets.
Revised: training set size and initialization scale, we train a depth-3 width-200 MLP with ReLU activations with the AdamW optimizer using MSE loss with one-hot targets.

Original: We find that the network quickly fits the training set, and test accuracy improves much later, as shown in Figure 3d, just as in the stereotypical grokking learning first observed in algorithmic datasets.
Revised: We find that the network quickly fits the training set, and test accuracy improves much later, as shown in Figure 3d, just as in the stereotypical grokking learning first observed in algorithmic datasets.

Original: Figure 3e shows the effect of training set size on time to generalization for MNIST.
Revised: Figure 3e shows the effect of training set size on time to generalization for MNIST.

Original: We find a result similar to what Power et al.
Revised: We find a result similar to what Power et al.

Original: (2022) observed, namely that generalization time increases rapidly once one approaches a certain critical data set size.
Revised: (2022) observed, namely that generalization time increases rapidly once one approaches a certain critical data set size.

Revised: The conclusions still hold for the cross entropy loss (see Appendix E), although with quantitatively milder effects.

Original: Sentiment analysis of text We look for grokking using LSTMs (Hochreiter and Schmidhuber, 1997) for IMDb dataset (Maas et al., 2011).
Revised: Sentiment analysis of text We look for grokking using LSTMs (Hochreiter and Schmidhuber, 1997) for IMDb dataset (Maas et al., 2011).

Original: Similar to Eq.
Revised: Similar to Eq.

Original: (3), we reduce training and test losses to depend on only the weight norm w and data size N .
Revised: (3), we reduce training and test losses to depend on only the weight norm w and data size N .

Original: We show the reduced training and test error in Figure 4 (a)(b).
Revised: We show the reduced training and test error in Figure 4 (a)(b).

Original: For large data size, e.g., the full dataset, training and test errors have similar "U" shapes 5 , so one cannot create grokking via the "LU" mechanism.
Revised: For large data size, e.g., the full dataset, training and test errors have similar "U" shapes 5 , so one cannot create grokking via the "LU" mechanism.

Original: For small data size, say 1k, however, the mismatch between training and test errors makes it possible to create grokking via large initializations.
Revised: For small data size, say 1k, however, the mismatch between training and test errors makes it possible to create grokking via large initializations.

Original: In Figure 4 (c), we initialize weights larger ( α = 6 ) with weight decay 1, overfitting is complete within 10 2 steps, but generalization does not start until around 10 3 steps.
Revised: In Figure 4 (c), we initialize weights larger ( α = 6 ) with weight decay 1, overfitting is complete within 10 2 steps, but generalization does not start until around 10 3 steps.

Original: Note that the generalization "jump" is not as sharp as on algorithmic datasets (Power et al., 2022) or MNIST, but at least generalization is delayed here.
Revised: Note that the generalization "jump" is not as sharp as on algorithmic datasets (Power et al., 2022) or MNIST, but at least generalization is delayed here.

Original: By contrast, if we use the standard initialization ( α = 1 ) with no weight decay, generalization happens early on during training, and does not improve much after overfitting.
Revised: By contrast, if we use the standard initialization ( α = 1 ) with no weight decay, generalization happens early on during training, and does not improve much after overfitting.

Original: Molecules We search for grokking using the graph convolutional neural network (GCNN) for QM dataset (Ramakrishnan et al., 2014).
Revised: We search for grokking using the graph convolutional neural network (GCNN) for QM dataset (Ramakrishnan et al., 2014).

Original: Similar to Eq.
Revised: Similar to Eq.

Original: (3), we define the reduced training/test losses, which are only dependent on weight norm w and data size N .
Revised: (3), we define the reduced training/test losses, which are only dependent on weight norm w and data size N .

Original: As shown in Figure 5(a)(b), when data size is large, training and test losses have similar "U" shapes, hence grokking is impossible via the "LU mechanism".
Revised: As shown in Figure 5(a)(b), when data size is large, training and test losses have similar "U" shapes, hence grokking is impossible via the "LU mechanism".

Original: When data size is small, training and test losses mismatch somewhere in the region α = w/w 0 > 1 , making grokking possible.
Revised: When data size is small, training and test losses mismatch somewhere in the region α = w/w 0 > 1 , making grokking possible.

Original: Indeed, shown in Figure 5(d), there is a sharp drop in test loss around 10 4 steps if initialization is 3 times larger than standard, while standard initialization does not lead to grokking.
Revised: Indeed, shown in Figure 5(d), there is a sharp drop in test loss around 10 4 steps if initialization is 3 times larger than standard, while standard initialization does not lead to grokking.

Original: Note that zero weight decay is applied in both cases, implying the existence of implicit regularizations.
Revised: Note that zero weight decay is applied in both cases, implying the existence of implicit regularizations.

Original:  R EPRESENTATION IS KEY TO GROKKING
Revised:  R EPRESENTATION IS KEY TO GROKKING

Original: In Section 4, we showed that increasing initialization scales can make grokking happen for standard
Revised: In Section 4, we showed that increasing initialization scales can make grokking happen for standard

Original: ML tasks.
Revised: ML tasks.

Original: However, this seems a bit artificial and does not explain why standard initialization leads to grokking on algorithmic datasets, but not on standard ML datasets, say MNIST.
Revised: However, this seems a bit artificial and does not explain why standard initialization leads to grokking on algorithmic datasets, but not on standard ML datasets, say MNIST.

Original: The key difference is how much the task relies on representation learning.
Revised: The key difference is how much the task relies on representation learning.

Original: For the MNIST dataset, the quality of representation determines whether the test accuracy is 95% or 100%; by contrast in algorithmic datasets, the quality of representation determines whether test accuracy is random guess (bad representation) or 100% (good representation).
Revised: For the MNIST dataset, the quality of representation determines whether the test accuracy is 95% or 100%; by contrast in algorithmic datasets, the quality of representation determines whether test accuracy is random guess (bad representation) or 100% (good representation).

Original: So overfitting (under a bad representation) has a more dramatic effect on algorithmic datasets, i.e., the model weights increase quickly during overfitting but test accuracy remains low.
Revised: So overfitting (under a bad representation) has a more dramatic effect on algorithmic datasets, i.e., the model weights increase quickly during overfitting but test accuracy remains low.

Original: During overfitting, model weight norm is much larger than at initialization, but then drops below the initialization norm when the model generalizes, shown in
Revised: During overfitting, model weight norm is much larger than at initialization, but then drops below the initialization norm when the model generalizes, shown in

Original: Figure 8 (see Appendix B), and also observed by Nanda and Lieberum (2022).
Revised: Figure 8 (see Appendix B), and also observed by Nanda et al.

Original: In the following, we will compare algorithmic datasets (Section 5.1) to MNIST (Section 5.2).
Revised: In the following, we will compare algorithmic datasets (Section 5.1) to MNIST (Section 5.2).

Original: We show how their loss landscapes depend on representations differently, and how the difference leads to different outcomes (grokking or not).
Revised: We show how their loss landscapes depend on representations differently, and how the difference leads to different outcomes (grokking or not).

Original: 5. A LGORITHMIC DATASETS
Revised: 5. A LGORITHMIC DATASETS

Original: Setup Algorithmic datasets are the task of learning a binary operation a ◦ b
Revised: Setup Algorithmic datasets are the task of learning a binary operation a ◦ b

Original: =
Revised: =

Original: c ( a, b, c are symbols) with neural networks, which aim to predict c from input ( a, b ) .
Revised: c ( a, b, c are symbols) with neural networks, which aim to predict c from input ( a, b ) .

Original: We take the toy addition setup in (Liu et al., 2022), where each input digit 0 ≤ i ≤ p − 1 (output label 0 ≤ k ≤ 2( q − 1) ) is embedded as a vector E
Revised: We take the toy addition setup in (Liu et al., 2022), where each input digit 0 ≤ i ≤ p − 1 (output label 0 ≤ k ≤ 2( q − 1) ) is embedded as a vector E

Original: i
Revised: i

Original: ( Y k ).
Revised: ( Y k ).

Original: A decoder MLP is employed to predict Y k
Revised: A decoder MLP is employed to predict Y k

Original: =
Revised: =

Original: Dec( E i + E j ) ( k = i + j ) .
Revised: Dec( E i + E j ) ( k = i + j ) .

Original: In the setup of grokking, both the decoder and the input representations R ≡ { E i } are trainable, with learningrates η D and η R , respectively; in the setup of landscape analysis, only decoder is trainable, as we [Figure] n t a t i
Revised: In the setup of grokking, both the decoder and the input representations R ≡ { E i } are trainable, with learning

Original: [Figure]

Original: m

Original: [Figure]
Revised: [Figure]

Original: ss i n e ss m

Original: I
Revised: I

Original: II
Revised: II

Original: [Figure]
Revised: [Figure]

Original: E

Original: [Figure]
Revised: [Figure]

Original:  r e p r e s e n t a t i o n v = v 0 = cos g h
Revised: good r e p r e s e n t a t i o n v = v 0 = cos g h

Original: L
Revised: L

Original: L
Revised: L

Original: L(f)  explain below.
Revised: L(f) rates η D and η R , respectively; in the setup of landscape analysis, only decoder is trainable, as we explain below.

Original: Training and test losses depend on three factors: (i) representation R , (ii) weight norm w and (iii) weight direction ˆ w .
Revised: Training and test losses depend on three factors: (i) representation R , (ii) weight norm w and (iii) weight direction ˆ w .

Original: As in previous sections, we can optimize ˆ w by minimizing the training loss on constant weight norm spheres.
Revised: As in previous sections, we can optimize ˆ w by minimizing the training loss on constant weight norm spheres.

Original: We further reduce the high-dimensional representations to 1D by interpolating in a particular direction:
Revised: We further reduce the high-dimensional representations to 1D by interpolating in a particular direction:

Original: [Equation]
Revised: [Equation]

Original: where R linear refers to the linear representation in which number k is embedded to E k
Revised: where R linear refers to the linear representation in which number k is embedded to E k

Original: =
Revised: =

Original: [ k, 0 , · · · , 0] ,
Revised: [ k, 0 , · · · , 0] ,

Original: R random is the initialized representation drawn from Gaussian distributions, i.e, E k ∼ N ( 0 , I ) , and m ∈ [0 , 1] is a scalar interpolating between R linear and R random , that we term representation messiness because R = R linear when m = 0 , and R = R random when m = 1 .
Revised: R random is the initialized representation drawn from Gaussian distributions, i.e, E k ∼ N ( 0 , I ) , and m ∈ [0 , 1] is a scalar interpolating between R linear and R random , that we term representation messiness because R = R linear when m = 0 , and R = R random when m = 1 .

Original: After these reductions, both training and test losses become functions of two variables, representation messiness m and weight norm w :
Revised: After these reductions, both training and test losses become functions of two variables, representation messiness m and weight norm w :

Original: [Equation]
Revised: [Equation]

Original: Note that our definition of ˜ l train ( w, m ) excludes the weight decay term ℓ reg = 12 γw 2 , but we should be aware of its presence when we analyze the dynamics of ( w, m ) , which is governed by the gradient flow on ˜ l train ( w, m ) plus weight decay ( η R /η
Revised: Note that our definition of ˜ l train ( w, m ) excludes the weight decay term ℓ reg = 12 γw 2 , but we should be aware of its presence when we analyze the dynamics of ( w, m ) , which is governed by the gradient flow on ˜ l train ( w, m ) plus weight decay ( η R /η

Original: D are learning rates of representation/decoder):
Revised: D are learning rates of representation/decoder):

Original: [Equation]
Revised: [Equation]

Original: More experimental details are included in Appendix D.
Revised: More experimental details are included in Appendix D.

Original: Landscape We show ˜ l train ( w, m ) and ˜ l test ( w, m ) in Figures 6a and 6b, indicating the generalizing solution with a green star.
Revised: Landscape We show ˜ l train ( w, m ) and ˜ l test ( w, m ) in Figures 6a and 6b, indicating the generalizing solution with a green star.

Original: Based on the reduced training loss (Figure 6a), we can divide the 2D plane into two regions I and II , separated by a dashed yellow line (the contour of training loss = 0.05): ( I ):
Revised: Based on the reduced training loss (Figure 6a), we can divide the 2D plane into two regions I and II , separated by a dashed yellow line (the contour of training loss = 0.05): ( I ):

Original: The darker region, with high training losses/gradients and small weight norm.
Revised: The darker region, with high training losses/gradients and small weight norm.

Original: ( II ): The lighter region, with low training losses/gradients and large weight norm.
Revised: ( II ): The lighter region, with low training losses/gradients and large weight norm.

Original: Comparing Figures 6a and 6b reveals that training and test loss landscapes differ, especially in region II .
Revised: Comparing Figures 6a and 6b reveals that training and test loss landscapes differ, especially in region II .

Original: Moreover, while the training loss depends weakly on m , the test loss depends strongly on m .
Revised: Moreover, while the training loss depends weakly on m , the test loss depends strongly on m .

Original: As we will see, the (weak) dependence of training loss on representation drives the model to the generalizing solution.
Revised: As we will see, the (weak) dependence of training loss on representation drives the model to the generalizing solution.

Original: However, the driving force is small because the dependence is weak, leading to grokking.
Revised: However, the driving force is small because the dependence is weak, leading to grokking.

Original: We elaborate below how these particular loss landscapes lead to grokking.
Revised: We elaborate below how these particular loss landscapes lead to grokking.

Original: Grokking dynamics In region II , the dynamics is slow (for small γ ) due to nearly vanishing gradients.
Revised: Grokking dynamics In region II , the dynamics is slow (for small γ ) due to nearly vanishing gradients.

Original: By contrast, the dynamics in region I is relatively fast.
Revised: By contrast, the dynamics in region I is relatively fast.

Original: As we will explain, dynamics is also slow on the boundary of I and II , and grokking is the consequence of traversing region II and/or the boundary.
Revised: As we will explain, dynamics is also slow on the boundary of I and II , and grokking is the consequence of traversing region II and/or the boundary.

Original: Let us analyze a typical path A to E shown in Figure 6(a)(b).
Revised: Let us analyze a typical path A to E shown in Figure 6(a)(b).

Original: A rolls "downhill" to B following training gradients, possibly continuing to C due to momentum.
Revised: A rolls "downhill" to B following training gradients, possibly continuing to C due to momentum.

Original: C is located in II where ˜ l train ≈ 0 , so according to Eq.
Revised: C is located in II where ˜ l train ≈ 0 , so according to Eq.

Original: (6), dm/dt ≈ 0 and dw/dt ≈ − η D γw or, equivalently, d (log w ) /dt
Revised: (6), dm/dt ≈ 0 and dw/dt ≈ − η D γw or, equivalently, d (log w ) /dt

Original: ≈ − η D γ .
Revised: ≈ − η D γ .

Original: So (log w, m ) moves with a constant speed v =
Revised: So (log w, m ) moves with a constant speed v =

Original: η D γ in the − w direction from C to D , a point near the boundary.
Revised: η D γ in the − w direction from C to D , a point near the boundary.

Original: Negative gradients around the boundary point towards larger w and smaller m , shown in Figure 6d (a zoom-in of Figure 6a).
Revised: Negative gradients around the boundary point towards larger w and smaller m , shown in Figure 6d (a zoom-in of Figure 6a).

Original: The gradients become increasingly large as the model goes deeper inside region I , and at some point, the gradient totally cancels out v in the gradient direction, making the model start to drift along the boundary, as illustrated in Figure 6f.
Revised: The gradients become increasingly large as the model goes deeper inside region I , and at some point, the gradient totally cancels out v in the gradient direction, making the model start to drift along the boundary, as illustrated in Figure 6f.

Original: Then the model moves along the boundary with a new velocity v ′
Revised: Then the model moves along the boundary with a new velocity v ′

Original: = v cos θ 6 , until it reaches the generalizing solution E .
Revised: = v cos θ 6 , until it reaches the generalizing solution E .

Original: The above picture is supported by empirical experiments in Appendix B and also Nanda and Lieberum (2022).
Revised: The above picture is supported by empirical experiments in Appendix B and also Nanda et al.

Original: Based on the picture, we also show the ability to eliminate grokking in Appendix B.
Revised: Based on the picture, we also show the ability to eliminate grokking in Appendix B.

Original: The slow dynamics from C to E is the origin of grokking.
Revised: The slow dynamics from C to E is the origin of grokking.

Original: During this period, the model first moves in the − w direction with a velocity v over the distance L 1 =
Revised: During this period, the model first moves in the − w direction with a velocity v over the distance L 1 =

Original: L − h cot θ , and then moves along the boundary with a velocity v ′ over the distance
Revised: L − h cot θ , and then moves along the boundary with a velocity v ′ over the distance

Original: L 2 = h/ sin θ .
Revised: L 2 = h/ sin θ .

Original: So the total time is t =
Revised: So the total time is t =

Original: L 1 /v + L 2 /v ′
Revised: L 1 /v + L 2 /v ′

Original: =
Revised: =

Original: ( L + h tan θ ) /
Revised: ( L + h tan θ ) /

Original: ( η D γ ) .
Revised: ( η D γ ) .

Original: This formula agrees with the observation that large weight decays γ and/or larger decoder learning rates η D can make generalization happen faster (Power et al., 2022;
Revised: This formula agrees with the observation that large weight decays γ and/or larger decoder learning rates η D can make generalization happen faster (Power et al., 2022;

Original: Liu et al., 2022).
Revised: Liu et al., 2022).

Original: Besides, the path manifests intriguing multiple descent of test loss, shown in Figure 6c.
Revised: Besides, the path manifests intriguing multiple descent of test loss, shown in Figure 6c.

Original: Dependence of grokking on training data size Another important observation in Power et al.
Revised: Dependence of grokking on training data size Another important observation in Power et al.

Original: (2022) is that grokking happens faster for larger training size.
Revised: (2022) is that grokking happens faster for larger training size.

Original: Our landscape analysis can also explain the data size dependence.
Revised: Our landscape analysis can also explain the data size dependence.

Original: In Figure 6e, we show the contours (training loss = 0.02) for different training sizes (25, 35, 45, 55).
Revised: In Figure 6e, we show the contours (training loss = 0.02) for different training sizes (25, 35, 45, 55).

Original: The contours of training size 45 and 55 both connect to the green star, meaning that generalization will eventually happen.
Revised: The contours of training size 45 and 55 both connect to the green star, meaning that generalization will eventually happen.

Original: However, the slopes of the contours are different, i.e., θ 55 < θ 45 .
Revised: However, the slopes of the contours are different, i.e., θ 55 < θ 45 .

Original: Since t =
Revised: Since t =

Original: ( L + h tan θ ) /
Revised: ( L + h tan θ ) /

Original: ( η D γ ) increases as θ increases, we have t 55 < t 45 , i.e, more training data leads to faster grokking.
Revised: ( η D γ ) increases as θ increases, we have t 55 < t 45 , i.e, more training data leads to faster grokking.

Original: For training size 35 and 25, the contours do not connect to the green star, so generalization will not happen, no matter how long the training will be run.
Revised: For training size 35 and 25, the contours do not connect to the green star, so generalization will not happen, no matter how long the training will be run.

Original: 5. MNIST
Revised: 5. MNIST

Original: We now study how training and test losses depend on representation messiness in the MNIST dataset.
Revised: We now study how training and test losses depend on representation messiness in the MNIST dataset.

Original: We denote the 28 × 28 images as the raw representation R raw .
Revised: We denote the 28 × 28 images as the raw representation R raw .

Original: We construct a linearly separable representation R linear by assigning input representations proportional to their label y
Revised: We construct a linearly separable representation R linear by assigning input representations proportional to their label y

Original: i , for example, an image of a 2 is represented by a matrix with all elements being 2.
Revised: i , for example, an image of a 2 is represented by a matrix with all elements being 2.

Original: Similar to Eq.
Revised: Similar to Eq.

Original: (4), we use m ∈ [0 , 1] to interpolated between R raw and R linear :
Revised: (4), we use m ∈ [0 , 1] to interpolated between R raw and R linear :

Original: [Equation]
Revised: [Equation]

Original: Similarly to Eq.
Revised: Similarly to Eq.

Original: (5), we define and plot ˜ l train ( w, m ) and ˜ l test ( w, m ) in Figure 7, using the full dataset
Revised: (5), we define and plot ˜ l train ( w, m ) and ˜ l test ( w, m ) in Figure 7, using the full dataset

Original: N = 60000 .
Revised: N = 60000 .

Original: Comparing Figures 7a and 7b reveals two things: (1) The training and test losses behave similarly; (2) Both training and test losses depend very weakly on m .
Revised: Comparing Figures 7a and 7b reveals two things: (1) The training and test losses behave similarly; (2) Both training and test losses depend very weakly on m .

Original: This implies that the raw image representation is already quite close to being optimal, so decent test accuracy can be obtained even without learning optimal representations.
Revised: This implies that the raw image representation is already quite close to being optimal, so decent test accuracy can be obtained even without learning optimal representations.

Original: As a result, grokking does not occur (Figure 7c).
Revised: As a result, grokking does not occur (Figure 7c).

Original: Comparing Figure 6 and 7, we see that the (strong) dependence of test performance on the representation is the key to grokking: the dependence on representation is strong for algorithmic datasets, so grokking happens.
Revised: Comparing Figure 6 and 7, we see that the (strong) dependence of test performance on the representation is the key to grokking: the dependence on representation is strong for algorithmic datasets, so grokking happens.

Original: By contrast, the dependence is weak for MNIST, so grokking does not happen.
Revised: By contrast, the dependence is weak for MNIST, so grokking does not happen.

Original: 10.0 relative weight norm = w / w 0.
Revised: 10.0 relative weight norm = w / w 0.

Original: 1.0 r e p r e s e n t a t i o n m e ss i n e ss m
Revised: 1.0 r e p r e s e n t a t i o n m e ss i n e ss m

Original: A
Revised: A

Original: B
Revised: B

Original: C training loss 0 .0 .0 .0 .0 .0.
Revised: C training loss 0 .0 .0 .0 .0 .0.

Original: (a) 0.
Revised: (a) 0.

Original: 10.0 relative weight norm = w / w 0.
Revised: 10.0 relative weight norm = w / w 0.

Original: 1.0 r e p r e s e n t a t i o n m e ss i n e ss m
Revised: 1.0 r e p r e s e n t a t i o n m e ss i n e ss m

Original: A
Revised: A

Original: B
Revised: B

Original: C test loss 0 .0 .0 .0 .0.
Revised: C test loss 0 .0 .0 .0 .0.

Original: (b)
Revised: (b)

Original: A
Revised: A

Original: B
Revised: B

Original: C Path 0.
Revised: C Path 0.

Original: l o ss train test (c)
Revised: l o ss train test (c)

Original:  R ELATION TO R ELATED WORKS
Revised:  R ELATION TO R ELATED WORKS

Original: Grokking was first observed for algorithmic datasets by Power et al.
Revised: Grokking was first observed for algorithmic datasets by Power et al.

Original: Several attempts have been made to understand grokking: (a) Liu et al.
Revised: Several attempts have been made to understand grokking: (a) Liu et al.

Original: (2022) attributes grokking to the slow formation of good representations.
Revised: (2022) attributes grokking to the slow formation of good representations.

Original: (b) Shah (2021) suggests that generalizable solutions achieve lower loss than overfitting solutions, providing a training signal encouraging generalization.
Revised: (b) Shah (2021) suggests that generalizable solutions achieve lower loss than overfitting solutions, providing a training signal encouraging generalization.

Original: (c) Nanda and
Revised: (c) Nanda et al.

Original: Lieberum (2022) suggests grokking is a phase change due to limited data and regularization.
Revised: (2023) suggests grokking is a phase change due to limited data and regularization.

Original: (d) Barak et al.
Revised: (d) Barak et al.

Original: (2022) suggests that generalization is due not to random search, but to hidden progress of SGD to gradually amplify a Fourier gap.
Revised: (2022) suggests that generalization is due not to random search, but to hidden progress of SGD to gradually amplify a Fourier gap.

Original: (e) Thilak et al.
Revised: (e) Thilak et al.

Original: (2022) links grokking to the "Slingshot mechanism" specific to adaptive optimizers.
Revised: (2022) links grokking to the "Slingshot mechanism" specific to adaptive optimizers.

Original: (f) Millidge (2022) describes training as a random walk over parameters.
Revised: (f) Millidge (2022) describes training as a random walk over parameters.

Original: Our conclusion supports (a)(b)(c)(d), but does not necessarily negate (e)(f).
Revised: Our conclusion supports (a)(b)(c)(d), but does not necessarily negate (e)(f).

Original: Double descent is the phenomenon that performance first gets worse and then gets better as we increase the model size, data size, training epochs or regularization (Nakkiran et al., 2021; Yilmaz and Heckel, 2022; Nakkiran, 2019).
Revised: Double descent is the phenomenon that performance first gets worse and then gets better as we increase the model size, data size, training epochs or regularization (Nakkiran et al., 2021; Yilmaz and Heckel, 2022; Nakkiran, 2019).

Original: The typical "U" shape of test loss in this paper does not conflict with double descent, because we are plotting the weight norm instead of the number of model parameters (Ng and Ma, 2022).
Revised: The typical "U" shape of test loss in this paper does not conflict with double descent, because we are plotting the weight norm instead of the number of model parameters (Ng and Ma, 2022).

Original: However, the "U"-shape should better be considered as empirically common rather than provably universal.
Revised: However, the "U"-shape should better be considered as empirically common rather than provably universal.

Original: In fact, the interaction between properties of data and inductive biases of learning algorithms can be more complicated than double descent (Chen et al., 2021; d’Ascoli et al., 2020).
Revised: In fact, the interaction between properties of data and inductive biases of learning algorithms can be more complicated than double descent (Chen et al., 2021; d’Ascoli et al., 2020).

Original: Initialization From the optimization perspective, initializations are usually based on the "edge of chaos" idea such that variance of features and gradients should be preserved in the forward and backward pass (Glorot and Bengio, 2010;
Revised: Initialization From the optimization perspective, initializations are usually based on the "edge of chaos" idea such that variance of features and gradients should be preserved in the forward and backward pass (Glorot and Bengio, 2010;

Original: He et al., 2015; Bahri et al., 2020; Yang and Schoenholz, 2017; Jing et al., 2017), or based on analyzing Jacobians and/or Hessians (Skorski et al., 2020).
Revised: He et al., 2015; Bahri et al., 2020; Yang and Schoenholz, 2017; Jing et al., 2017), or based on analyzing Jacobians and/or Hessians (Skorski et al., 2020).

Original: From the generalization perspective, it was shown that large initializations overfit data easily but result in poor generalization (Xu et al., 2019; Zhang et al., 2020), which agrees with our LU mechanism.
Revised: From the generalization perspective, it was shown that large initializations overfit data easily but result in poor generalization (Xu et al., 2019; Zhang et al., 2020), which agrees with our LU mechanism.

Original: Weight decay regularization is a standard trick in machine learning and has various effects on optimization and generalization (Zhang et al., 2018; Van Laarhoven, 2017).
Revised: Weight decay regularization is a standard trick in machine learning and has various effects on optimization and generalization (Zhang et al., 2018; Van Laarhoven, 2017).

Original: In particular, Lewkowycz and Gur-Ari (2020) observes that it takes t ∝ 1 /λ training steps to reach maximum test performance.
Revised: In particular, Lewkowycz and Gur-Ari (2020) observes that it takes t ∝ 1 /λ training steps to reach maximum test performance.

Original: This is strikingly similar to the grokking time t ∝ 1 /λ we derived from the LU mechanism.
Revised: This is strikingly similar to the grokking time t ∝ 1 /λ we derived from the LU mechanism.

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: This study elucidates the grokking phenomenon from the perspective of loss landscapes.
Revised: This study elucidates the grokking phenomenon from the perspective of loss landscapes.

Original: Our conclusions are: (i) grokking originates from the mismatch between training and test losses  ("LU" mechanism).
Revised: Our conclusions are: (i) grokking originates from the mismatch between training and test losses at high model weight norm ("LU" mechanism).

Original: (ii) grokking can happen in various models for a wide range of datasets, although the grokking signature is usually most dramatic for algorithmic datasets.
Revised: (ii) grokking can happen in various models for a wide range of datasets, although the grokking signature is usually most dramatic for algorithmic datasets.

Original: (iii)
Revised: (iii)

Original: The severity of grokking depends on how much the task relies on learning representations.
Revised: The severity of grokking depends on how much the task relies on learning representations.

Original: This work not only reveals the mechanism of grokking, but also shows that reduced landscape analysis is a useful tool for characterizing data-model interaction and representation learning.
Revised: This work not only reveals the mechanism of grokking, but also shows that reduced landscape analysis is a useful tool for characterizing data-model interaction and representation learning.

Original: Appendix A E XPERIMENT DETAILS
Revised: Appendix A E XPERIMENT DETAILS

Original: Sentiment analysis of text IMDb (Maas et al., 2011) includes 50k movie reviews to be classified as being positive or negative.
Revised: Sentiment analysis of text IMDb (Maas et al., 2011) includes 50k movie reviews to be classified as being positive or negative.

Original: To pre-process the data, we extract the 1000 most frequent words and tokenize each review into an array of token indices.
Revised: To pre-process the data, we extract the 1000 most frequent words and tokenize each review into an array of token indices.

Original: Less frequent words are ignored, and each review array is padded to length 500.
Revised: Less frequent words are ignored, and each review array is padded to length 500.

Original: We adopt the LSTM model (Hochreiter and Schmidhuber, 1997) to perform the classification, with two layers, embedding dimension 64, and hidden dimension 128.
Revised: We adopt the LSTM model (Hochreiter and Schmidhuber, 1997) to perform the classification, with two layers, embedding dimension 64, and hidden dimension 128.

Original: We use the Adam optimizer (Kingma and Ba, 2014) with learning rate 0.001 to minimize the binary cross entropy loss.
Revised: We use the Adam optimizer (Kingma and Ba, 2014) with learning rate 0.001 to minimize the binary cross entropy loss.

Original: We hold back 25% of the dataset for testing.
Revised: We hold back 25% of the dataset for testing.

Original: Molecules QM9 is a database for small molecules and their properties.
Revised: Molecules QM9 is a database for small molecules and their properties.

Original: We use a graph convolutional neural network (GCNN) to predict the isotropic polarizability.
Revised: We use a graph convolutional neural network (GCNN) to predict the isotropic polarizability.

Original: The GCNN contains 2 convolutional layers with ReLU activation, followed by a linear layer.
Revised: The GCNN contains 2 convolutional layers with ReLU activation, followed by a linear layer.

Original: We use the Adam optimizer with learning rate 0.001 to minimize the MSE loss.
Revised: We use the Adam optimizer with learning rate 0.001 to minimize the MSE loss.

Original: We split the dataset into 50/50 train/test.
Revised: We split the dataset into 50/50 train/test.

Original: MNIST We train width-200 depth-3 ReLU MLPs on the MNIST dataset with MSE loss.
Revised: MNIST We train width-200 depth-3 ReLU MLPs on the MNIST dataset with MSE loss.

Original: We use the
Revised: We use the

Original: AdamW optimizer with a learning rate of 0.001 and a batch size of 200.
Revised: AdamW optimizer with a learning rate of 0.001 and a batch size of 200.

Original: B W EIGHT N ORM E VOLUTION O VER T IME ON A LGORITHMIC T ASKS
Revised: B W EIGHT N ORM E VOLUTION O VER T IME ON A LGORITHMIC T ASKS

Original: Evolution of weight norm As mentioned in Section 5, the dynamics of model weight norm over the course of training, on algorithmic tasks, support the LU mechanism picture of grokking.
Revised: Evolution of weight norm As mentioned in Section 5, the dynamics of model weight norm over the course of training, on algorithmic tasks, support the LU mechanism picture of grokking.

Original: Figure 8a, shows how model norm changes over time and we see that there is an initial increase in weight norm, which peaks during overfitting, but then drops during the period of generalization to be lower than the initialization norm.
Revised: Figure 8a, shows how model norm changes over time and we see that there is an initial increase in weight norm, which peaks during overfitting, but then drops during the period of generalization to be lower than the initialization norm.

Original: For this experiment, we used the setup of (Nanda and Lieberum, 2022), training a 1-layer Transformer on modular addition ( p = 113 ).
Revised: For this experiment, we used the setup of (Nanda et al., 2023), training a 1-layer Transformer on modular addition ( p = 113 ).

Original: The model width d model = 128 , with 4 attention heads, and d mlp = 512 with ReLU activations.
Revised: The model width d model = 128 , with 4 attention heads, and d mlp = 512 with ReLU activations.

Original: We train with AdamW with a learning rate of 0.
Revised: We train with AdamW with a learning rate of 0.

Original: and weight decay γ = 1 .
Revised: and weight decay γ = 1 .

Original: Constraining a small weight norm eliminates grokking As shown in Figure 8b, reducing the initialization scale ( α = 0 .
Revised: Constraining a small weight norm eliminates grokking As shown in Figure 8b, reducing the initialization scale ( α = 0 .

Original: 8) and constraining optimization to hold model weight norm constant over training brings train accuracy and test accuracy learning curves together, almost eliminating grokking.
Revised: 8) and constraining optimization to hold model weight norm constant over training brings train accuracy and test accuracy learning curves together, almost eliminating grokking.

Original: We would like to investigate in future works whether this training trick can be helpful for more standard machine learning tasks
Revised: We would like to investigate in future works whether this training trick can be helpful for more standard machine learning tasks.

Original: [Figure]
Revised: [Figure]

Original: C TIME TO GENERALIZE VERSUS WEIGHT DECAY
Revised: C TIME TO GENERALIZE VERSUS WEIGHT DECAY

Original: In our discussion of the “LU mechanism” as an explanation for grokking in Section 2, we predicted that the training time required for a model to generalize should be t ∝ γ − 1 where γ is the weight decay.
Revised: In our discussion of the “LU mechanism” as an explanation for grokking in Section 2, we predicted that the training time required for a model to generalize should be t ∝ γ − 1 where γ is the weight decay.

Original: To test this, we perform a grid search over weight decays γ and plot the number of training steps required for models to reach a specified level of test accuracy in Figure 9a-9b.
Revised: To test this, we perform a grid search over weight decays γ and plot the number of training steps required for models to reach a specified level of test accuracy in Figure 9a-9b.

Original: We also show full training curves for these runs in Figure 9c-9d.
Revised: We also show full training curves for these runs in Figure 9c-9d.

Original: We perform experiments in two setups:
Revised: We perform experiments in two setups:

Original: (a) Transformer on modular addition : We use the replication of grokking from Nanda and Lieberum (2022) and train a 1-layer transformer on modular addition ( p = 113 and a train set fraction of 0 . 3 ) where d model = 128 , with 4 attention heads, d mlp = 512 , ReLU activations, and an AdamW learning rate of 0.001.
Revised: (a) Transformer on modular addition : We use the replication of grokking from Nanda et al. and train a 1-layer transformer on modular addition ( p = 113 and a train set fraction of 0 . 3 ) where d model = 128 , with 4 attention heads, d mlp = 512 , ReLU activations, and an AdamW learning rate of 0.001.

Original: From Figure 9a, we find that t ∝ γ − 1 holds across roughly two orders of magnitude of t and γ .
Revised: From Figure 9a, we find that t ∝ γ − 1 holds across roughly two orders of magnitude of t and γ .

Original: There is some seed dependence on the generalization time (some seeds consistently require longer to generalize), but for each seed (corresponding to a particular model initialization) the relation t ∝ γ − 1 appears to fit the data well.
Revised: There is some seed dependence on the generalization time (some seeds consistently require longer to generalize), but for each seed (corresponding to a particular model initialization) the relation t ∝ γ − 1 appears to fit the data well.

Original: (b) ReLU MLP on MNIST : We train ReLU MLPs on MNIST as described in Appendix A.
Revised: (b) ReLU MLP on MNIST : We train ReLU MLPs on MNIST as described in Appendix A.

Original: We use an α = 9 . 0 and train on a reduced training set of 1000 samples to delay generalization / induce grokking.
Revised: We use an α = 9 . 0 and train on a reduced training set of 1000 samples to delay generalization / induce grokking.

Original: From Figure 9b, we find that for γ roughly between 0.1 and 1.0 the relation t ∝ γ − 1 holds.
Revised: From Figure 9b, we find that for γ roughly between 0.1 and 1.0 the relation t ∝ γ − 1 holds.

Original: Very high values of weight decay seem to mess with optimization.
Revised: Very high values of weight decay seem to mess with optimization.

Original: On the other hand, with very low weight decay the model generalizes faster than naively expected, perhaps due to implicit regularization.
Revised: On the other hand, with very low weight decay the model generalizes faster than naively expected, perhaps due to implicit regularization [Figure] t (a) [Figure]

Revised: i

Original: [Figure]
Revised: [Figure]

Revised: i z a t i

Revised: [Figure]

Revised: n

Revised: [Figure]

Revised: p

Revised: [Figure]

Revised: T e

Revised: [Figure]

Revised: cc

Revised: [Figure]

Revised: a

Revised: [Figure]

Original: D S ECTION 5.1 SETUP
Revised: t (b) [Figure] [Figure] (d)

Original: [Table]
Revised: [Table]

Revised: D S ECTION 5.1 SETUP

Original: Architecture Similar to Liu et al.
Revised: Architecture Similar to Liu et al.

Original: (2022), the decoder architecture is an MLP with hard coded addition.
Revised: (2022), the decoder architecture is an MLP with hard coded addition.

Original: Each input symbol i is encoded to a scalar E i .
Revised: Each input symbol i is encoded to a scalar E i .

Original: Each output symbol k is represented by a 30D random vector ˆ Y k .
Revised: Each output symbol k is represented by a 30D random vector ˆ Y k .

Original: We consider addition with base p , so input 0 ≤ i, j ≤ p − 1 and output 0 ≤ k
Revised: We consider addition with base p , so input 0 ≤ i, j ≤ p − 1 and output 0 ≤ k

Original: = i + j ≤ 2( p − 1) .
Revised: = i + j ≤ 2( p − 1) .

Original: We denote representation as R =
Revised: We denote representation as R =

Original: { E 0 , E 1 · · · , E p − 1 } .
Revised: { E 0 , E 1 · · · , E p − 1 } .

Original: The MLP has two hidden layers, with neurons 1-200-200-30 in each layer and ReLU activations.
Revised: The MLP has two hidden layers, with neurons 1-200-200-30 in each layer and ReLU activations.

Original: Given a training sample ( E i , E j )
Revised: Given a training sample ( E i , E j )

Original: → Y k where i + j = k , the prediction of the MLP decoder is
Revised: → Y k where i + j = k , the prediction of the MLP decoder is

Original: [Equation]
Revised: [Equation]

Original: and the loss function being the mean squared error (MSE) between Y k and ˆ Y k , and w being the decoder weight.
Revised: and the loss function being the mean squared error (MSE) between Y k and ˆ Y k , and w being the decoder weight.

Original: Although the common setup of grokking is to make both the representation R and the decoder w trainable, we will freeze part of them for easier analysis.
Revised: Although the common setup of grokking is to make both the representation R and the decoder w trainable, we will freeze part of them for easier analysis.

Original: This is where it could be a bit confusing, so we explicitly distinguish three setups: landscape analysis , reduced trajectory analysis and full trajectory analysis .
Revised: This is where it could be a bit confusing, so we explicitly distinguish three setups: landscape analysis , reduced trajectory analysis and full trajectory analysis .

Original: Each setup have different subset of trainable parameters, as shown in
Revised: Each setup have different subset of trainable parameters, as shown in

Original: Table 1.
Revised: Table 1.

Original: Landscape analysis Both the representation R and weight norm w are fixed.
Revised: Landscape analysis Both the representation R and weight norm w are fixed.

Original: Only the weight direction ˆ w is trainable.
Revised: Only the weight direction ˆ w is trainable.

Original: The representation R is fixed according to Eq.
Revised: The representation R is fixed according to Eq.

Original: (4), which is dependent on m , the representation messiness.
Revised: (4), which is dependent on m , the representation messiness.

Original: The decoder has fixed weight norm w , but the weight direction ˆ w is trainable.
Revised: The decoder has fixed weight norm w , but the weight direction ˆ w is trainable.

Original: For each fixed ( w, m ) , we minimize training loss over ˆ w to get
Revised: For each fixed ( w, m ) , we minimize training loss over ˆ w to get

Original: [Equation]
Revised: [Equation]

Original: and define reduced training and test loss, as in Eq.
Revised: and define reduced training and test loss, as in Eq.

Original: The minimization is implemented by the
Revised: The minimization is implemented by the

Original: Adam optimizer with learning rate 10 − 3 for 10 4 steps.
Revised: Adam optimizer with learning rate 10 − 3 for 10 4 steps.

Original: Although ( w, m ) are not trainable, we repeat the above minimization independently for different ( w, m ) .
Revised: Although ( w, m ) are not trainable, we repeat the above minimization independently for different ( w, m ) .

Original: In Figure 6 (a)(b)(d), the background heatmaps belong to landscape analysis.
Revised: In Figure 6 (a)(b)(d), the background heatmaps belong to landscape analysis.

Original: Reduced trajectory analysis is a “thought experiment" based on landscape analysis.
Revised: Reduced trajectory analysis is a “thought experiment" based on landscape analysis.

Original: Since full trajectory analysis can be intractable due to too high dimensions, we try to reduce the trajectory anaysis to 2D, by making two assumptions about the real dynamics: (1) Scale separation : the dynamics of ˆ w is much faster than the dynamics along w and along m , such that ˆ w ( t )
Revised: Since full trajectory analysis can be intractable due to too high dimensions, we try to reduce the trajectory anaysis to 2D, by making two assumptions about the real dynamics: (1) Scale separation : the dynamics of ˆ w is much faster than the dynamics along w and along m , such that ˆ w ( t )

Original: =
Revised: =

Original: ˆ w ∗ ( w ( t ) , m ( t )) is valid at every moment during training.
Revised: ˆ w ∗ ( w ( t ) , m ( t )) is valid at every moment during training.

Original: (2) Representation evolution is linear , i.e., interpolating between initial random Gaussian and final linear representation.
Revised: (2) Representation evolution is linear , i.e., interpolating between initial random Gaussian and final linear representation.

Original: With these two assumptions, the training dynamics is effectively reduced to 2D, depending only on ( w, m ) , obeying Eq.
Revised: With these two assumptions, the training dynamics is effectively reduced to 2D, depending only on ( w, m ) , obeying Eq.

Original: In Figure(a)(b)(c), the path A → E belongs to reduced trajectory analysis.
Revised: In Figure(a)(b)(c), the path A → E belongs to reduced trajectory analysis.

Original: Admittedly the reduced trajectory may deviate from the full trajectory since the assumptions may not be met, but it can shed light on the full trajectory: the weight norm first increases and then increases, and the decrease of weight norm is highly correlated with generalization (please see Appendix B and
Revised: Admittedly the reduced trajectory may deviate from the full trajectory since the assumptions may not be met, but it can shed light on the full trajectory: the weight norm first increases and then increases, and the decrease of weight norm is highly correlated with generalization (please see Appendix B and

Original: Figure 8.
Revised: Figure 8.

Revised: E MNIST EXPERIMENTS WITH CROSS ENTROPY LOSS

Revised: To respond to a reviewer’s concern that our use of the MSE loss is the “secret" to get grokking on MNIST (Figure 3), we reran our experiments with the cross entropy (CE) loss.

Revised: The results are qualitatively similar, with some quantitative differences.

Revised: Landscape analysis Comparing Figure 3 (MSE) and Figure 10 (CE), we notice the they are qualitatively similar: (1) for small datasets, the reduced training error and test error resemble an “L" and

Revised: [Figure]

Revised: [Figure]

Revised: “U" against the weight norm, respectively; (2) for large datasets, the “U" becomes more like “L", i e., the mismatch between the reduced training and test error is small.

Revised: However, a quantitative difference exist: CE produces a broader “Goldilocks zone" (the weight range where generalization happens) than MSE.

Revised: This implies that to induce grokking with CE, we need to increase the weight norm to a larger value (say α = 100 )

Revised: . Training dynamics

Revised: We are able to observe delayed generalization during trianing on MNIST with cross entropy loss, but doing so requires a higher α than was necessary when using MSE loss, as predicted by the reduced loss landscapes in Figure 10.

Revised: Figure 11 shows training trajectories from a 3-layer ReLU MLP on

Revised: MNIST trained with cross entropy loss with α = 100 and D = 200 .

Revised: We see that test accuracy rises to 30-40% early in training, then plateaus for an extended period, before increasing to ≈ 75% while train accuracy remains at 100%.

Revised: While the dynamics are not as clean as with MSE loss, since test accuracy first plateaus at better-than-random accuracy, we think it is still fair to classify these dynamics as “grokking” due to the improvement in generalization late in training after a plateau.

Revised: [Figure]